1 | use alloc::vec::Vec; |
---|---|

2 | use std::fmt; |

3 | use std::iter::once; |

4 | |

5 | use super::lazy_buffer::LazyBuffer; |

6 | |

7 | /// An iterator adaptor that iterates through all the `k`-permutations of the |

8 | /// elements from an iterator. |

9 | /// |

10 | /// See [`.permutations()`](crate::Itertools::permutations) for |

11 | /// more information. |

12 | #[must_use= "iterator adaptors are lazy and do nothing unless consumed"] |

13 | pub struct Permutations<I: Iterator> { |

14 | vals: LazyBuffer<I>, |

15 | state: PermutationState, |

16 | } |

17 | |

18 | impl<I> Clone for Permutations<I> |

19 | where I: Clone + Iterator, |

20 | I::Item: Clone, |

21 | { |

22 | clone_fields!(vals, state); |

23 | } |

24 | |

25 | #[derive(Clone, Debug)] |

26 | enum PermutationState { |

27 | StartUnknownLen { |

28 | k: usize, |

29 | }, |

30 | OngoingUnknownLen { |

31 | k: usize, |

32 | min_n: usize, |

33 | }, |

34 | Complete(CompleteState), |

35 | Empty, |

36 | } |

37 | |

38 | #[derive(Clone, Debug)] |

39 | enum CompleteState { |

40 | Start { |

41 | n: usize, |

42 | k: usize, |

43 | }, |

44 | Ongoing { |

45 | indices: Vec<usize>, |

46 | cycles: Vec<usize>, |

47 | } |

48 | } |

49 | |

50 | enum CompleteStateRemaining { |

51 | Known(usize), |

52 | Overflow, |

53 | } |

54 | |

55 | impl<I> fmt::Debug for Permutations<I> |

56 | where I: Iterator + fmt::Debug, |

57 | I::Item: fmt::Debug, |

58 | { |

59 | debug_fmt_fields!(Permutations, vals, state); |

60 | } |

61 | |

62 | pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> { |

63 | let mut vals = LazyBuffer::new(iter); |

64 | |

65 | if k == 0 { |

66 | // Special case, yields single empty vec; `n` is irrelevant |

67 | let state = PermutationState::Complete(CompleteState::Start { n: 0, k: 0 }); |

68 | |

69 | return Permutations { |

70 | vals, |

71 | state |

72 | }; |

73 | } |

74 | |

75 | let mut enough_vals = true; |

76 | |

77 | while vals.len() < k { |

78 | if !vals.get_next() { |

79 | enough_vals = false; |

80 | break; |

81 | } |

82 | } |

83 | |

84 | let state = if enough_vals { |

85 | PermutationState::StartUnknownLen { k } |

86 | } else { |

87 | PermutationState::Empty |

88 | }; |

89 | |

90 | Permutations { |

91 | vals, |

92 | state |

93 | } |

94 | } |

95 | |

96 | impl<I> Iterator for Permutations<I> |

97 | where |

98 | I: Iterator, |

99 | I::Item: Clone |

100 | { |

101 | type Item = Vec<I::Item>; |

102 | |

103 | fn next(&mut self) -> Option<Self::Item> { |

104 | self.advance(); |

105 | |

106 | let &mut Permutations { ref vals, ref state } = self; |

107 | |

108 | match *state { |

109 | PermutationState::StartUnknownLen { .. } => panic!("unexpected iterator state"), |

110 | PermutationState::OngoingUnknownLen { k, min_n } => { |

111 | let latest_idx = min_n - 1; |

112 | let indices = (0..(k - 1)).chain(once(latest_idx)); |

113 | |

114 | Some(indices.map(|i| vals[i].clone()).collect()) |

115 | } |

116 | PermutationState::Complete(CompleteState::Ongoing { ref indices, ref cycles }) => { |

117 | let k = cycles.len(); |

118 | Some(indices[0..k].iter().map(|&i| vals[i].clone()).collect()) |

119 | }, |

120 | PermutationState::Complete(CompleteState::Start { .. }) | PermutationState::Empty => None |

121 | } |

122 | } |

123 | |

124 | fn count(self) -> usize { |

125 | fn from_complete(complete_state: CompleteState) -> usize { |

126 | match complete_state.remaining() { |

127 | CompleteStateRemaining::Known(count) => count, |

128 | CompleteStateRemaining::Overflow => { |

129 | panic!("Iterator count greater than usize::MAX"); |

130 | } |

131 | } |

132 | } |

133 | |

134 | let Permutations { vals, state } = self; |

135 | match state { |

136 | PermutationState::StartUnknownLen { k } => { |

137 | let n = vals.len() + vals.it.count(); |

138 | let complete_state = CompleteState::Start { n, k }; |

139 | |

140 | from_complete(complete_state) |

141 | } |

142 | PermutationState::OngoingUnknownLen { k, min_n } => { |

143 | let prev_iteration_count = min_n - k + 1; |

144 | let n = vals.len() + vals.it.count(); |

145 | let complete_state = CompleteState::Start { n, k }; |

146 | |

147 | from_complete(complete_state) - prev_iteration_count |

148 | }, |

149 | PermutationState::Complete(state) => from_complete(state), |

150 | PermutationState::Empty => 0 |

151 | } |

152 | } |

153 | |

154 | fn size_hint(&self) -> (usize, Option<usize>) { |

155 | match self.state { |

156 | PermutationState::StartUnknownLen { .. } | |

157 | PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound? |

158 | PermutationState::Complete(ref state) => match state.remaining() { |

159 | CompleteStateRemaining::Known(count) => (count, Some(count)), |

160 | CompleteStateRemaining::Overflow => (::std::usize::MAX, None) |

161 | } |

162 | PermutationState::Empty => (0, Some(0)) |

163 | } |

164 | } |

165 | } |

166 | |

167 | impl<I> Permutations<I> |

168 | where |

169 | I: Iterator, |

170 | I::Item: Clone |

171 | { |

172 | fn advance(&mut self) { |

173 | let &mut Permutations { ref mut vals, ref mut state } = self; |

174 | |

175 | *state = match *state { |

176 | PermutationState::StartUnknownLen { k } => { |

177 | PermutationState::OngoingUnknownLen { k, min_n: k } |

178 | } |

179 | PermutationState::OngoingUnknownLen { k, min_n } => { |

180 | if vals.get_next() { |

181 | PermutationState::OngoingUnknownLen { k, min_n: min_n + 1 } |

182 | } else { |

183 | let n = min_n; |

184 | let prev_iteration_count = n - k + 1; |

185 | let mut complete_state = CompleteState::Start { n, k }; |

186 | |

187 | // Advance the complete-state iterator to the correct point |

188 | for _ in 0..(prev_iteration_count + 1) { |

189 | complete_state.advance(); |

190 | } |

191 | |

192 | PermutationState::Complete(complete_state) |

193 | } |

194 | } |

195 | PermutationState::Complete(ref mut state) => { |

196 | state.advance(); |

197 | |

198 | return; |

199 | } |

200 | PermutationState::Empty => { return; } |

201 | }; |

202 | } |

203 | } |

204 | |

205 | impl CompleteState { |

206 | fn advance(&mut self) { |

207 | *self = match *self { |

208 | CompleteState::Start { n, k } => { |

209 | let indices = (0..n).collect(); |

210 | let cycles = ((n - k)..n).rev().collect(); |

211 | |

212 | CompleteState::Ongoing { |

213 | cycles, |

214 | indices |

215 | } |

216 | }, |

217 | CompleteState::Ongoing { ref mut indices, ref mut cycles } => { |

218 | let n = indices.len(); |

219 | let k = cycles.len(); |

220 | |

221 | for i in (0..k).rev() { |

222 | if cycles[i] == 0 { |

223 | cycles[i] = n - i - 1; |

224 | |

225 | let to_push = indices.remove(i); |

226 | indices.push(to_push); |

227 | } else { |

228 | let swap_index = n - cycles[i]; |

229 | indices.swap(i, swap_index); |

230 | |

231 | cycles[i] -= 1; |

232 | return; |

233 | } |

234 | } |

235 | |

236 | CompleteState::Start { n, k } |

237 | } |

238 | } |

239 | } |

240 | |

241 | fn remaining(&self) -> CompleteStateRemaining { |

242 | use self::CompleteStateRemaining::{Known, Overflow}; |

243 | |

244 | match *self { |

245 | CompleteState::Start { n, k } => { |

246 | if n < k { |

247 | return Known(0); |

248 | } |

249 | |

250 | let count: Option<usize> = (n - k + 1..n + 1).fold(Some(1), |acc, i| { |

251 | acc.and_then(|acc| acc.checked_mul(i)) |

252 | }); |

253 | |

254 | match count { |

255 | Some(count) => Known(count), |

256 | None => Overflow |

257 | } |

258 | } |

259 | CompleteState::Ongoing { ref indices, ref cycles } => { |

260 | let mut count: usize = 0; |

261 | |

262 | for (i, &c) in cycles.iter().enumerate() { |

263 | let radix = indices.len() - i; |

264 | let next_count = count.checked_mul(radix) |

265 | .and_then(|count| count.checked_add(c)); |

266 | |

267 | count = match next_count { |

268 | Some(count) => count, |

269 | None => { return Overflow; } |

270 | }; |

271 | } |

272 | |

273 | Known(count) |

274 | } |

275 | } |

276 | } |

277 | } |

278 |