itertools/
permutations.rs

1use alloc::boxed::Box;
2use alloc::vec::Vec;
3use std::fmt;
4use std::iter::once;
5use std::iter::FusedIterator;
6
7use super::lazy_buffer::LazyBuffer;
8use crate::size_hint::{self, SizeHint};
9
10/// An iterator adaptor that iterates through all the `k`-permutations of the
11/// elements from an iterator.
12///
13/// See [`.permutations()`](crate::Itertools::permutations) for
14/// more information.
15#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
16pub struct Permutations<I: Iterator> {
17    vals: LazyBuffer<I>,
18    state: PermutationState,
19}
20
21impl<I> Clone for Permutations<I>
22where
23    I: Clone + Iterator,
24    I::Item: Clone,
25{
26    clone_fields!(vals, state);
27}
28
29#[derive(Clone, Debug)]
30enum PermutationState {
31    /// No permutation generated yet.
32    Start { k: usize },
33    /// Values from the iterator are not fully loaded yet so `n` is still unknown.
34    Buffered { k: usize, min_n: usize },
35    /// All values from the iterator are known so `n` is known.
36    Loaded {
37        indices: Box<[usize]>,
38        cycles: Box<[usize]>,
39    },
40    /// No permutation left to generate.
41    End,
42}
43
44impl<I> fmt::Debug for Permutations<I>
45where
46    I: Iterator + fmt::Debug,
47    I::Item: fmt::Debug,
48{
49    debug_fmt_fields!(Permutations, vals, state);
50}
51
52pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> {
53    Permutations {
54        vals: LazyBuffer::new(iter),
55        state: PermutationState::Start { k },
56    }
57}
58
59impl<I> Iterator for Permutations<I>
60where
61    I: Iterator,
62    I::Item: Clone,
63{
64    type Item = Vec<I::Item>;
65
66    fn next(&mut self) -> Option<Self::Item> {
67        let Self { vals, state } = self;
68        match state {
69            PermutationState::Start { k: 0 } => {
70                *state = PermutationState::End;
71                Some(Vec::new())
72            }
73            &mut PermutationState::Start { k } => {
74                vals.prefill(k);
75                if vals.len() != k {
76                    *state = PermutationState::End;
77                    return None;
78                }
79                *state = PermutationState::Buffered { k, min_n: k };
80                Some(vals[0..k].to_vec())
81            }
82            PermutationState::Buffered { ref k, min_n } => {
83                if vals.get_next() {
84                    let item = (0..*k - 1)
85                        .chain(once(*min_n))
86                        .map(|i| vals[i].clone())
87                        .collect();
88                    *min_n += 1;
89                    Some(item)
90                } else {
91                    let n = *min_n;
92                    let prev_iteration_count = n - *k + 1;
93                    let mut indices: Box<[_]> = (0..n).collect();
94                    let mut cycles: Box<[_]> = (n - k..n).rev().collect();
95                    // Advance the state to the correct point.
96                    for _ in 0..prev_iteration_count {
97                        if advance(&mut indices, &mut cycles) {
98                            *state = PermutationState::End;
99                            return None;
100                        }
101                    }
102                    let item = vals.get_at(&indices[0..*k]);
103                    *state = PermutationState::Loaded { indices, cycles };
104                    Some(item)
105                }
106            }
107            PermutationState::Loaded { indices, cycles } => {
108                if advance(indices, cycles) {
109                    *state = PermutationState::End;
110                    return None;
111                }
112                let k = cycles.len();
113                Some(vals.get_at(&indices[0..k]))
114            }
115            PermutationState::End => None,
116        }
117    }
118
119    fn count(self) -> usize {
120        let Self { vals, state } = self;
121        let n = vals.count();
122        state.size_hint_for(n).1.unwrap()
123    }
124
125    fn size_hint(&self) -> SizeHint {
126        let (mut low, mut upp) = self.vals.size_hint();
127        low = self.state.size_hint_for(low).0;
128        upp = upp.and_then(|n| self.state.size_hint_for(n).1);
129        (low, upp)
130    }
131}
132
133impl<I> FusedIterator for Permutations<I>
134where
135    I: Iterator,
136    I::Item: Clone,
137{
138}
139
140fn advance(indices: &mut [usize], cycles: &mut [usize]) -> bool {
141    let n = indices.len();
142    let k = cycles.len();
143    // NOTE: if `cycles` are only zeros, then we reached the last permutation.
144    for i in (0..k).rev() {
145        if cycles[i] == 0 {
146            cycles[i] = n - i - 1;
147            indices[i..].rotate_left(1);
148        } else {
149            let swap_index = n - cycles[i];
150            indices.swap(i, swap_index);
151            cycles[i] -= 1;
152            return false;
153        }
154    }
155    true
156}
157
158impl PermutationState {
159    fn size_hint_for(&self, n: usize) -> SizeHint {
160        // At the beginning, there are `n!/(n-k)!` items to come.
161        let at_start = |n, k| {
162            debug_assert!(n >= k);
163            let total = (n - k + 1..=n).try_fold(1usize, |acc, i| acc.checked_mul(i));
164            (total.unwrap_or(usize::MAX), total)
165        };
166        match *self {
167            Self::Start { k } if n < k => (0, Some(0)),
168            Self::Start { k } => at_start(n, k),
169            Self::Buffered { k, min_n } => {
170                // Same as `Start` minus the previously generated items.
171                size_hint::sub_scalar(at_start(n, k), min_n - k + 1)
172            }
173            Self::Loaded {
174                ref indices,
175                ref cycles,
176            } => {
177                let count = cycles.iter().enumerate().try_fold(0usize, |acc, (i, &c)| {
178                    acc.checked_mul(indices.len() - i)
179                        .and_then(|count| count.checked_add(c))
180                });
181                (count.unwrap_or(usize::MAX), count)
182            }
183            Self::End => (0, Some(0)),
184        }
185    }
186}
itertools/permutations.rs

itertools/
permutations.rs
