itertools/
combinations.rs

1use alloc::boxed::Box;
2use core::array;
3use core::borrow::BorrowMut;
4use std::fmt;
5use std::iter::FusedIterator;
6
7use super::lazy_buffer::LazyBuffer;
8use alloc::vec::Vec;
9
10use crate::adaptors::checked_binomial;
11
12/// Iterator for `Vec` valued combinations returned by [`.combinations()`](crate::Itertools::combinations)
13pub type Combinations<I> = CombinationsGeneric<I, Vec<usize>>;
14/// Iterator for const generic combinations returned by [`.array_combinations()`](crate::Itertools::array_combinations)
15pub type ArrayCombinations<I, const K: usize> = CombinationsGeneric<I, [usize; K]>;
16
17/// Create a new `Combinations` from a cloneable iterator.
18pub fn combinations<I: Iterator>(iter: I, k: usize) -> Combinations<I>
19where
20    I::Item: Clone,
21{
22    Combinations::new(iter, (0..k).collect())
23}
24
25/// Create a new `ArrayCombinations` from a cloneable iterator.
26pub fn array_combinations<I: Iterator, const K: usize>(iter: I) -> ArrayCombinations<I, K>
27where
28    I::Item: Clone,
29{
30    ArrayCombinations::new(iter, array::from_fn(|i| i))
31}
32
33/// An iterator to iterate through all the `k`-length combinations in an iterator.
34///
35/// See [`.combinations()`](crate::Itertools::combinations) and [`.array_combinations()`](crate::Itertools::array_combinations) for more information.
36#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
37pub struct CombinationsGeneric<I: Iterator, Idx> {
38    indices: Idx,
39    pool: LazyBuffer<I>,
40    first: bool,
41}
42
43/// A type holding indices of elements in a pool or buffer of items from an inner iterator
44/// and used to pick out different combinations in a generic way.
45pub trait PoolIndex<T>: BorrowMut<[usize]> {
46    type Item;
47
48    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Self::Item
49    where
50        T: Clone;
51
52    fn len(&self) -> usize {
53        self.borrow().len()
54    }
55}
56impl<T> PoolIndex<T> for Box<[usize]> {
57    type Item = Vec<T>;
58
59    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
60    where
61        T: Clone,
62    {
63        pool.get_at(self)
64    }
65}
66impl<T> PoolIndex<T> for Vec<usize> {
67    type Item = Vec<T>;
68
69    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
70    where
71        T: Clone,
72    {
73        pool.get_at(self)
74    }
75}
76
77impl<T, const K: usize> PoolIndex<T> for [usize; K] {
78    type Item = [T; K];
79
80    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> [T; K]
81    where
82        T: Clone,
83    {
84        pool.get_array(*self)
85    }
86}
87
88impl<I, Idx> Clone for CombinationsGeneric<I, Idx>
89where
90    I: Iterator + Clone,
91    I::Item: Clone,
92    Idx: Clone,
93{
94    clone_fields!(indices, pool, first);
95}
96
97impl<I, Idx> fmt::Debug for CombinationsGeneric<I, Idx>
98where
99    I: Iterator + fmt::Debug,
100    I::Item: fmt::Debug,
101    Idx: fmt::Debug,
102{
103    debug_fmt_fields!(Combinations, indices, pool, first);
104}
105
106impl<I: Iterator, Idx: PoolIndex<I::Item>> CombinationsGeneric<I, Idx> {
107    /// Constructor with arguments the inner iterator and the initial state for the indices.
108    fn new(iter: I, indices: Idx) -> Self {
109        Self {
110            indices,
111            pool: LazyBuffer::new(iter),
112            first: true,
113        }
114    }
115
116    /// Returns the length of a combination produced by this iterator.
117    #[inline]
118    pub fn k(&self) -> usize {
119        self.indices.len()
120    }
121
122    /// Returns the (current) length of the pool from which combination elements are
123    /// selected. This value can change between invocations of [`next`](Combinations::next).
124    #[inline]
125    pub fn n(&self) -> usize {
126        self.pool.len()
127    }
128
129    /// Returns a reference to the source pool.
130    #[inline]
131    pub(crate) fn src(&self) -> &LazyBuffer<I> {
132        &self.pool
133    }
134
135    /// Return the length of the inner iterator and the count of remaining combinations.
136    pub(crate) fn n_and_count(self) -> (usize, usize) {
137        let Self {
138            indices,
139            pool,
140            first,
141        } = self;
142        let n = pool.count();
143        (n, remaining_for(n, first, indices.borrow()).unwrap())
144    }
145
146    /// Initialises the iterator by filling a buffer with elements from the
147    /// iterator. Returns true if there are no combinations, false otherwise.
148    fn init(&mut self) -> bool {
149        self.pool.prefill(self.k());
150        let done = self.k() > self.n();
151        if !done {
152            self.first = false;
153        }
154
155        done
156    }
157
158    /// Increments indices representing the combination to advance to the next
159    /// (in lexicographic order by increasing sequence) combination. For example
160    /// if we have n=4 & k=2 then `[0, 1] -> [0, 2] -> [0, 3] -> [1, 2] -> ...`
161    ///
162    /// Returns true if we've run out of combinations, false otherwise.
163    fn increment_indices(&mut self) -> bool {
164        // Borrow once instead of noise each time it's indexed
165        let indices = self.indices.borrow_mut();
166
167        if indices.is_empty() {
168            return true; // Done
169        }
170        // Scan from the end, looking for an index to increment
171        let mut i: usize = indices.len() - 1;
172
173        // Check if we need to consume more from the iterator
174        if indices[i] == self.pool.len() - 1 {
175            self.pool.get_next(); // may change pool size
176        }
177
178        while indices[i] == i + self.pool.len() - indices.len() {
179            if i > 0 {
180                i -= 1;
181            } else {
182                // Reached the last combination
183                return true;
184            }
185        }
186
187        // Increment index, and reset the ones to its right
188        indices[i] += 1;
189        for j in i + 1..indices.len() {
190            indices[j] = indices[j - 1] + 1;
191        }
192        // If we've made it this far, we haven't run out of combos
193        false
194    }
195
196    /// Returns the n-th item or the number of successful steps.
197    pub(crate) fn try_nth(&mut self, n: usize) -> Result<<Self as Iterator>::Item, usize>
198    where
199        I: Iterator,
200        I::Item: Clone,
201    {
202        let done = if self.first {
203            self.init()
204        } else {
205            self.increment_indices()
206        };
207        if done {
208            return Err(0);
209        }
210        for i in 0..n {
211            if self.increment_indices() {
212                return Err(i + 1);
213            }
214        }
215        Ok(self.indices.extract_item(&self.pool))
216    }
217}
218
219impl<I, Idx> Iterator for CombinationsGeneric<I, Idx>
220where
221    I: Iterator,
222    I::Item: Clone,
223    Idx: PoolIndex<I::Item>,
224{
225    type Item = Idx::Item;
226    fn next(&mut self) -> Option<Self::Item> {
227        let done = if self.first {
228            self.init()
229        } else {
230            self.increment_indices()
231        };
232
233        if done {
234            return None;
235        }
236
237        Some(self.indices.extract_item(&self.pool))
238    }
239
240    fn nth(&mut self, n: usize) -> Option<Self::Item> {
241        self.try_nth(n).ok()
242    }
243
244    fn size_hint(&self) -> (usize, Option<usize>) {
245        let (mut low, mut upp) = self.pool.size_hint();
246        low = remaining_for(low, self.first, self.indices.borrow()).unwrap_or(usize::MAX);
247        upp = upp.and_then(|upp| remaining_for(upp, self.first, self.indices.borrow()));
248        (low, upp)
249    }
250
251    #[inline]
252    fn count(self) -> usize {
253        self.n_and_count().1
254    }
255}
256
257impl<I, Idx> FusedIterator for CombinationsGeneric<I, Idx>
258where
259    I: Iterator,
260    I::Item: Clone,
261    Idx: PoolIndex<I::Item>,
262{
263}
264
265impl<I: Iterator> Combinations<I> {
266    /// Resets this `Combinations` back to an initial state for combinations of length
267    /// `k` over the same pool data source. If `k` is larger than the current length
268    /// of the data pool an attempt is made to prefill the pool so that it holds `k`
269    /// elements.
270    pub(crate) fn reset(&mut self, k: usize) {
271        self.first = true;
272
273        if k < self.indices.len() {
274            self.indices.truncate(k);
275            for i in 0..k {
276                self.indices[i] = i;
277            }
278        } else {
279            for i in 0..self.indices.len() {
280                self.indices[i] = i;
281            }
282            self.indices.extend(self.indices.len()..k);
283            self.pool.prefill(k);
284        }
285    }
286}
287
288/// For a given size `n`, return the count of remaining combinations or None if it would overflow.
289fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
290    let k = indices.len();
291    if n < k {
292        Some(0)
293    } else if first {
294        checked_binomial(n, k)
295    } else {
296        // https://en.wikipedia.org/wiki/Combinatorial_number_system
297        // http://www.site.uottawa.ca/~lucia/courses/5165-09/GenCombObj.pdf
298
299        // The combinations generated after the current one can be counted by counting as follows:
300        // - The subsequent combinations that differ in indices[0]:
301        //   If subsequent combinations differ in indices[0], then their value for indices[0]
302        //   must be at least 1 greater than the current indices[0].
303        //   As indices is strictly monotonically sorted, this means we can effectively choose k values
304        //   from (n - 1 - indices[0]), leading to binomial(n - 1 - indices[0], k) possibilities.
305        // - The subsequent combinations with same indices[0], but differing indices[1]:
306        //   Here we can choose k - 1 values from (n - 1 - indices[1]) values,
307        //   leading to binomial(n - 1 - indices[1], k - 1) possibilities.
308        // - (...)
309        // - The subsequent combinations with same indices[0..=i], but differing indices[i]:
310        //   Here we can choose k - i values from (n - 1 - indices[i]) values: binomial(n - 1 - indices[i], k - i).
311        //   Since subsequent combinations can in any index, we must sum up the aforementioned binomial coefficients.
312
313        // Below, `n0` resembles indices[i].
314        indices.iter().enumerate().try_fold(0usize, |sum, (i, n0)| {
315            sum.checked_add(checked_binomial(n - 1 - *n0, k - i)?)
316        })
317    }
318}
itertools/combinations.rs

itertools/
combinations.rs
