/rust/registry/src/index.crates.io-1949cf8c6b5b557f/itertools-0.10.5/src/permutations.rs
Line | Count | Source |
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 | 0 | pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> { |
63 | 0 | let mut vals = LazyBuffer::new(iter); |
64 | | |
65 | 0 | if k == 0 { |
66 | | // Special case, yields single empty vec; `n` is irrelevant |
67 | 0 | let state = PermutationState::Complete(CompleteState::Start { n: 0, k: 0 }); |
68 | | |
69 | 0 | return Permutations { |
70 | 0 | vals, |
71 | 0 | state |
72 | 0 | }; |
73 | 0 | } |
74 | | |
75 | 0 | let mut enough_vals = true; |
76 | | |
77 | 0 | while vals.len() < k { |
78 | 0 | if !vals.get_next() { |
79 | 0 | enough_vals = false; |
80 | 0 | break; |
81 | 0 | } |
82 | | } |
83 | | |
84 | 0 | let state = if enough_vals { |
85 | 0 | PermutationState::StartUnknownLen { k } |
86 | | } else { |
87 | 0 | PermutationState::Empty |
88 | | }; |
89 | | |
90 | 0 | Permutations { |
91 | 0 | vals, |
92 | 0 | state |
93 | 0 | } |
94 | 0 | } |
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 | 0 | fn next(&mut self) -> Option<Self::Item> { |
104 | 0 | self.advance(); |
105 | | |
106 | 0 | let &mut Permutations { ref vals, ref state } = self; |
107 | | |
108 | 0 | match *state { |
109 | 0 | PermutationState::StartUnknownLen { .. } => panic!("unexpected iterator state"), |
110 | 0 | PermutationState::OngoingUnknownLen { k, min_n } => { |
111 | 0 | let latest_idx = min_n - 1; |
112 | 0 | let indices = (0..(k - 1)).chain(once(latest_idx)); |
113 | | |
114 | 0 | Some(indices.map(|i| vals[i].clone()).collect()) |
115 | | } |
116 | 0 | PermutationState::Complete(CompleteState::Ongoing { ref indices, ref cycles }) => { |
117 | 0 | let k = cycles.len(); |
118 | 0 | Some(indices[0..k].iter().map(|&i| vals[i].clone()).collect()) |
119 | | }, |
120 | 0 | PermutationState::Complete(CompleteState::Start { .. }) | PermutationState::Empty => None |
121 | | } |
122 | 0 | } |
123 | | |
124 | 0 | fn count(self) -> usize { |
125 | 0 | fn from_complete(complete_state: CompleteState) -> usize { |
126 | 0 | match complete_state.remaining() { |
127 | 0 | CompleteStateRemaining::Known(count) => count, |
128 | | CompleteStateRemaining::Overflow => { |
129 | 0 | panic!("Iterator count greater than usize::MAX"); |
130 | | } |
131 | | } |
132 | 0 | } |
133 | | |
134 | 0 | let Permutations { vals, state } = self; |
135 | 0 | match state { |
136 | 0 | PermutationState::StartUnknownLen { k } => { |
137 | 0 | let n = vals.len() + vals.it.count(); |
138 | 0 | let complete_state = CompleteState::Start { n, k }; |
139 | | |
140 | 0 | from_complete(complete_state) |
141 | | } |
142 | 0 | PermutationState::OngoingUnknownLen { k, min_n } => { |
143 | 0 | let prev_iteration_count = min_n - k + 1; |
144 | 0 | let n = vals.len() + vals.it.count(); |
145 | 0 | let complete_state = CompleteState::Start { n, k }; |
146 | | |
147 | 0 | from_complete(complete_state) - prev_iteration_count |
148 | | }, |
149 | 0 | PermutationState::Complete(state) => from_complete(state), |
150 | 0 | PermutationState::Empty => 0 |
151 | | } |
152 | 0 | } |
153 | | |
154 | 0 | fn size_hint(&self) -> (usize, Option<usize>) { |
155 | 0 | match self.state { |
156 | | PermutationState::StartUnknownLen { .. } | |
157 | 0 | PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound? |
158 | 0 | PermutationState::Complete(ref state) => match state.remaining() { |
159 | 0 | CompleteStateRemaining::Known(count) => (count, Some(count)), |
160 | 0 | CompleteStateRemaining::Overflow => (::std::usize::MAX, None) |
161 | | } |
162 | 0 | PermutationState::Empty => (0, Some(0)) |
163 | | } |
164 | 0 | } |
165 | | } |
166 | | |
167 | | impl<I> Permutations<I> |
168 | | where |
169 | | I: Iterator, |
170 | | I::Item: Clone |
171 | | { |
172 | 0 | fn advance(&mut self) { |
173 | 0 | let &mut Permutations { ref mut vals, ref mut state } = self; |
174 | | |
175 | 0 | *state = match *state { |
176 | 0 | PermutationState::StartUnknownLen { k } => { |
177 | 0 | PermutationState::OngoingUnknownLen { k, min_n: k } |
178 | | } |
179 | 0 | PermutationState::OngoingUnknownLen { k, min_n } => { |
180 | 0 | if vals.get_next() { |
181 | 0 | PermutationState::OngoingUnknownLen { k, min_n: min_n + 1 } |
182 | | } else { |
183 | 0 | let n = min_n; |
184 | 0 | let prev_iteration_count = n - k + 1; |
185 | 0 | let mut complete_state = CompleteState::Start { n, k }; |
186 | | |
187 | | // Advance the complete-state iterator to the correct point |
188 | 0 | for _ in 0..(prev_iteration_count + 1) { |
189 | 0 | complete_state.advance(); |
190 | 0 | } |
191 | | |
192 | 0 | PermutationState::Complete(complete_state) |
193 | | } |
194 | | } |
195 | 0 | PermutationState::Complete(ref mut state) => { |
196 | 0 | state.advance(); |
197 | | |
198 | 0 | return; |
199 | | } |
200 | 0 | PermutationState::Empty => { return; } |
201 | | }; |
202 | 0 | } |
203 | | } |
204 | | |
205 | | impl CompleteState { |
206 | 0 | fn advance(&mut self) { |
207 | 0 | *self = match *self { |
208 | 0 | CompleteState::Start { n, k } => { |
209 | 0 | let indices = (0..n).collect(); |
210 | 0 | let cycles = ((n - k)..n).rev().collect(); |
211 | | |
212 | 0 | CompleteState::Ongoing { |
213 | 0 | cycles, |
214 | 0 | indices |
215 | 0 | } |
216 | | }, |
217 | 0 | CompleteState::Ongoing { ref mut indices, ref mut cycles } => { |
218 | 0 | let n = indices.len(); |
219 | 0 | let k = cycles.len(); |
220 | | |
221 | 0 | for i in (0..k).rev() { |
222 | 0 | if cycles[i] == 0 { |
223 | 0 | cycles[i] = n - i - 1; |
224 | 0 |
|
225 | 0 | let to_push = indices.remove(i); |
226 | 0 | indices.push(to_push); |
227 | 0 | } else { |
228 | 0 | let swap_index = n - cycles[i]; |
229 | 0 | indices.swap(i, swap_index); |
230 | | |
231 | 0 | cycles[i] -= 1; |
232 | 0 | return; |
233 | | } |
234 | | } |
235 | | |
236 | 0 | CompleteState::Start { n, k } |
237 | | } |
238 | | } |
239 | 0 | } |
240 | | |
241 | 0 | fn remaining(&self) -> CompleteStateRemaining { |
242 | | use self::CompleteStateRemaining::{Known, Overflow}; |
243 | | |
244 | 0 | match *self { |
245 | 0 | CompleteState::Start { n, k } => { |
246 | 0 | if n < k { |
247 | 0 | return Known(0); |
248 | 0 | } |
249 | | |
250 | 0 | let count: Option<usize> = (n - k + 1..n + 1).fold(Some(1), |acc, i| { |
251 | 0 | acc.and_then(|acc| acc.checked_mul(i)) |
252 | 0 | }); |
253 | | |
254 | 0 | match count { |
255 | 0 | Some(count) => Known(count), |
256 | 0 | None => Overflow |
257 | | } |
258 | | } |
259 | 0 | CompleteState::Ongoing { ref indices, ref cycles } => { |
260 | 0 | let mut count: usize = 0; |
261 | | |
262 | 0 | for (i, &c) in cycles.iter().enumerate() { |
263 | 0 | let radix = indices.len() - i; |
264 | 0 | let next_count = count.checked_mul(radix) |
265 | 0 | .and_then(|count| count.checked_add(c)); |
266 | | |
267 | 0 | count = match next_count { |
268 | 0 | Some(count) => count, |
269 | 0 | None => { return Overflow; } |
270 | | }; |
271 | | } |
272 | | |
273 | 0 | Known(count) |
274 | | } |
275 | | } |
276 | 0 | } |
277 | | } |