itertools/
combinations.rs

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

itertools/
combinations.rs
