/rust/registry/src/index.crates.io-1949cf8c6b5b557f/itertools-0.14.0/src/combinations.rs

Source
use core::array;
use core::borrow::BorrowMut;
use std::fmt;
use std::iter::FusedIterator;

use super::lazy_buffer::LazyBuffer;
use alloc::vec::Vec;

use crate::adaptors::checked_binomial;

/// Iterator for `Vec` valued combinations returned by [`.combinations()`](crate::Itertools::combinations)
pub type Combinations<I> = CombinationsGeneric<I, Vec<usize>>;
/// Iterator for const generic combinations returned by [`.array_combinations()`](crate::Itertools::array_combinations)
pub type ArrayCombinations<I, const K: usize> = CombinationsGeneric<I, [usize; K]>;

/// Create a new `Combinations` from a clonable iterator.
pub fn combinations<I: Iterator>(iter: I, k: usize) -> Combinations<I>
where
    I::Item: Clone,
{
    Combinations::new(iter, (0..k).collect())
}

/// Create a new `ArrayCombinations` from a clonable iterator.
pub fn array_combinations<I: Iterator, const K: usize>(iter: I) -> ArrayCombinations<I, K>
where
    I::Item: Clone,
{
    ArrayCombinations::new(iter, array::from_fn(|i| i))
}

/// An iterator to iterate through all the `k`-length combinations in an iterator.
///
/// See [`.combinations()`](crate::Itertools::combinations) and [`.array_combinations()`](crate::Itertools::array_combinations) for more information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct CombinationsGeneric<I: Iterator, Idx> {
    indices: Idx,
    pool: LazyBuffer<I>,
    first: bool,
}

/// A type holding indices of elements in a pool or buffer of items from an inner iterator
/// and used to pick out different combinations in a generic way.
pub trait PoolIndex<T>: BorrowMut<[usize]> {
    type Item;

    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Self::Item
    where
        T: Clone;

    fn len(&self) -> usize {
        self.borrow().len()
    }
}

impl<T> PoolIndex<T> for Vec<usize> {
    type Item = Vec<T>;

    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
    where
        T: Clone,
    {
        pool.get_at(self)
    }
}

impl<T, const K: usize> PoolIndex<T> for [usize; K] {
    type Item = [T; K];

    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> [T; K]
    where
        T: Clone,
    {
        pool.get_array(*self)
    }
}

impl<I, Idx> Clone for CombinationsGeneric<I, Idx>
where
    I: Iterator + Clone,
    I::Item: Clone,
    Idx: Clone,
{
    clone_fields!(indices, pool, first);
}

impl<I, Idx> fmt::Debug for CombinationsGeneric<I, Idx>
where
    I: Iterator + fmt::Debug,
    I::Item: fmt::Debug,
    Idx: fmt::Debug,
{
    debug_fmt_fields!(Combinations, indices, pool, first);
}

impl<I: Iterator, Idx: PoolIndex<I::Item>> CombinationsGeneric<I, Idx> {
    /// Constructor with arguments the inner iterator and the initial state for the indices.
    fn new(iter: I, indices: Idx) -> Self {
        Self {
            indices,
            pool: LazyBuffer::new(iter),
            first: true,
        }
    }

    /// Returns the length of a combination produced by this iterator.
    #[inline]
    pub fn k(&self) -> usize {
        self.indices.len()
    }

    /// Returns the (current) length of the pool from which combination elements are
    /// selected. This value can change between invocations of [`next`](Combinations::next).
    #[inline]
    pub fn n(&self) -> usize {
        self.pool.len()
    }

    /// Returns a reference to the source pool.
    #[inline]
    pub(crate) fn src(&self) -> &LazyBuffer<I> {
        &self.pool
    }

    /// Return the length of the inner iterator and the count of remaining combinations.
    pub(crate) fn n_and_count(self) -> (usize, usize) {
        let Self {
            indices,
            pool,
            first,
        } = self;
        let n = pool.count();
        (n, remaining_for(n, first, indices.borrow()).unwrap())
    }

    /// Initialises the iterator by filling a buffer with elements from the
    /// iterator. Returns true if there are no combinations, false otherwise.
    fn init(&mut self) -> bool {
        self.pool.prefill(self.k());
        let done = self.k() > self.n();
        if !done {
            self.first = false;
        }

        done
    }

    /// Increments indices representing the combination to advance to the next
    /// (in lexicographic order by increasing sequence) combination. For example
    /// if we have n=4 & k=2 then `[0, 1] -> [0, 2] -> [0, 3] -> [1, 2] -> ...`
    ///
    /// Returns true if we've run out of combinations, false otherwise.
    fn increment_indices(&mut self) -> bool {
        // Borrow once instead of noise each time it's indexed
        let indices = self.indices.borrow_mut();

        if indices.is_empty() {
            return true; // Done
        }
        // Scan from the end, looking for an index to increment
        let mut i: usize = indices.len() - 1;

        // Check if we need to consume more from the iterator
        if indices[i] == self.pool.len() - 1 {
            self.pool.get_next(); // may change pool size
        }

        while indices[i] == i + self.pool.len() - indices.len() {
            if i > 0 {
                i -= 1;
            } else {
                // Reached the last combination
                return true;
            }
        }

        // Increment index, and reset the ones to its right
        indices[i] += 1;
        for j in i + 1..indices.len() {
            indices[j] = indices[j - 1] + 1;
        }
        // If we've made it this far, we haven't run out of combos
        false
    }

    /// Returns the n-th item or the number of successful steps.
    pub(crate) fn try_nth(&mut self, n: usize) -> Result<<Self as Iterator>::Item, usize>
    where
        I: Iterator,
        I::Item: Clone,
    {
        let done = if self.first {
            self.init()
        } else {
            self.increment_indices()
        };
        if done {
            return Err(0);
        }
        for i in 0..n {
            if self.increment_indices() {
                return Err(i + 1);
            }
        }
        Ok(self.indices.extract_item(&self.pool))
    }
}

impl<I, Idx> Iterator for CombinationsGeneric<I, Idx>
where
    I: Iterator,
    I::Item: Clone,
    Idx: PoolIndex<I::Item>,
{
    type Item = Idx::Item;
    fn next(&mut self) -> Option<Self::Item> {
        let done = if self.first {
            self.init()
        } else {
            self.increment_indices()
        };

        if done {
            return None;
        }

        Some(self.indices.extract_item(&self.pool))
    }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.try_nth(n).ok()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (mut low, mut upp) = self.pool.size_hint();
        low = remaining_for(low, self.first, self.indices.borrow()).unwrap_or(usize::MAX);
        upp = upp.and_then(|upp| remaining_for(upp, self.first, self.indices.borrow()));
        (low, upp)
    }

    #[inline]
    fn count(self) -> usize {
        self.n_and_count().1
    }
}

impl<I, Idx> FusedIterator for CombinationsGeneric<I, Idx>
where
    I: Iterator,
    I::Item: Clone,
    Idx: PoolIndex<I::Item>,
{
}

impl<I: Iterator> Combinations<I> {
    /// Resets this `Combinations` back to an initial state for combinations of length
    /// `k` over the same pool data source. If `k` is larger than the current length
    /// of the data pool an attempt is made to prefill the pool so that it holds `k`
    /// elements.
    pub(crate) fn reset(&mut self, k: usize) {
        self.first = true;

        if k < self.indices.len() {
            self.indices.truncate(k);
            for i in 0..k {
                self.indices[i] = i;
            }
        } else {
            for i in 0..self.indices.len() {
                self.indices[i] = i;
            }
            self.indices.extend(self.indices.len()..k);
            self.pool.prefill(k);
        }
    }
}

/// For a given size `n`, return the count of remaining combinations or None if it would overflow.
fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
    let k = indices.len();
    if n < k {
        Some(0)
    } else if first {
        checked_binomial(n, k)
    } else {
        // https://en.wikipedia.org/wiki/Combinatorial_number_system
        // http://www.site.uottawa.ca/~lucia/courses/5165-09/GenCombObj.pdf

        // The combinations generated after the current one can be counted by counting as follows:
        // - The subsequent combinations that differ in indices[0]:
        //   If subsequent combinations differ in indices[0], then their value for indices[0]
        //   must be at least 1 greater than the current indices[0].
        //   As indices is strictly monotonically sorted, this means we can effectively choose k values
        //   from (n - 1 - indices[0]), leading to binomial(n - 1 - indices[0], k) possibilities.
        // - The subsequent combinations with same indices[0], but differing indices[1]:
        //   Here we can choose k - 1 values from (n - 1 - indices[1]) values,
        //   leading to binomial(n - 1 - indices[1], k - 1) possibilities.
        // - (...)
        // - The subsequent combinations with same indices[0..=i], but differing indices[i]:
        //   Here we can choose k - i values from (n - 1 - indices[i]) values: binomial(n - 1 - indices[i], k - i).
        //   Since subsequent combinations can in any index, we must sum up the aforementioned binomial coefficients.

        // Below, `n0` resembles indices[i].
        indices.iter().enumerate().try_fold(0usize, |sum, (i, n0)| {
            sum.checked_add(checked_binomial(n - 1 - *n0, k - i)?)
        })
    }
}

Coverage Report

Created: 2025-12-20 06:48

Line	Count	Source
1		use core::array;
2		use core::borrow::BorrowMut;
3		use std::fmt;
4		use std::iter::FusedIterator;
5
6		use super::lazy_buffer::LazyBuffer;
7		use alloc::vec::Vec;
8
9		use crate::adaptors::checked_binomial;
10
11		/// Iterator for `Vec` valued combinations returned by [`.combinations()`](crate::Itertools::combinations)
12		pub type Combinations<I> = CombinationsGeneric<I, Vec<usize>>;
13		/// Iterator for const generic combinations returned by [`.array_combinations()`](crate::Itertools::array_combinations)
14		pub type ArrayCombinations<I, const K: usize> = CombinationsGeneric<I, [usize; K]>;
15
16		/// Create a new `Combinations` from a clonable iterator.
17	0	pub fn combinations<I: Iterator>(iter: I, k: usize) -> Combinations<I>
18	0	where
19	0	I::Item: Clone,
20		{
21	0	Combinations::new(iter, (0..k).collect())
22	0	}
23
24		/// Create a new `ArrayCombinations` from a clonable iterator.
25	0	pub fn array_combinations<I: Iterator, const K: usize>(iter: I) -> ArrayCombinations<I, K>
26	0	where
27	0	I::Item: Clone,
28		{
29	0	ArrayCombinations::new(iter, array::from_fn(\|i\| i))
30	0	}
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"]
36		pub 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.
44		pub 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	0	fn len(&self) -> usize {
52	0	self.borrow().len()
53	0	}
54		}
55
56		impl<T> PoolIndex<T> for Vec<usize> {
57		type Item = Vec<T>;
58
59	0	fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
60	0	where
61	0	T: Clone,
62		{
63	0	pool.get_at(self)
64	0	}
65		}
66
67		impl<T, const K: usize> PoolIndex<T> for [usize; K] {
68		type Item = [T; K];
69
70	0	fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> [T; K]
71	0	where
72	0	T: Clone,
73		{
74	0	pool.get_array(*self)
75	0	}
76		}
77
78		impl<I, Idx> Clone for CombinationsGeneric<I, Idx>
79		where
80		I: Iterator + Clone,
81		I::Item: Clone,
82		Idx: Clone,
83		{
84		clone_fields!(indices, pool, first);
85		}
86
87		impl<I, Idx> fmt::Debug for CombinationsGeneric<I, Idx>
88		where
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
96		impl<I: Iterator, Idx: PoolIndex<I::Item>> CombinationsGeneric<I, Idx> {
97		/// Constructor with arguments the inner iterator and the initial state for the indices.
98	0	fn new(iter: I, indices: Idx) -> Self {
99	0	Self {
100	0	indices,
101	0	pool: LazyBuffer::new(iter),
102	0	first: true,
103	0	}
104	0	}
105
106		/// Returns the length of a combination produced by this iterator.
107		#[inline]
108	0	pub fn k(&self) -> usize {
109	0	self.indices.len()
110	0	}
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	0	pub fn n(&self) -> usize {
116	0	self.pool.len()
117	0	}
118
119		/// Returns a reference to the source pool.
120		#[inline]
121	0	pub(crate) fn src(&self) -> &LazyBuffer<I> {
122	0	&self.pool
123	0	}
124
125		/// Return the length of the inner iterator and the count of remaining combinations.
126	0	pub(crate) fn n_and_count(self) -> (usize, usize) {
127		let Self {
128	0	indices,
129	0	pool,
130	0	first,
131	0	} = self;
132	0	let n = pool.count();
133	0	(n, remaining_for(n, first, indices.borrow()).unwrap())
134	0	}
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	0	fn init(&mut self) -> bool {
139	0	self.pool.prefill(self.k());
140	0	let done = self.k() > self.n();
141	0	if !done {
142	0	self.first = false;
143	0	}
144
145	0	done
146	0	}
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	0	fn increment_indices(&mut self) -> bool {
154		// Borrow once instead of noise each time it's indexed
155	0	let indices = self.indices.borrow_mut();
156
157	0	if indices.is_empty() {
158	0	return true; // Done
159	0	}
160		// Scan from the end, looking for an index to increment
161	0	let mut i: usize = indices.len() - 1;
162
163		// Check if we need to consume more from the iterator
164	0	if indices[i] == self.pool.len() - 1 {
165	0	self.pool.get_next(); // may change pool size
166	0	}
167
168	0	while indices[i] == i + self.pool.len() - indices.len() {
169	0	if i > 0 {
170	0	i -= 1;
171	0	} else {
172		// Reached the last combination
173	0	return true;
174		}
175		}
176
177		// Increment index, and reset the ones to its right
178	0	indices[i] += 1;
179	0	for j in i + 1..indices.len() {
180	0	indices[j] = indices[j - 1] + 1;
181	0	}
182		// If we've made it this far, we haven't run out of combos
183	0	false
184	0	}
185
186		/// Returns the n-th item or the number of successful steps.
187	0	pub(crate) fn try_nth(&mut self, n: usize) -> Result<<Self as Iterator>::Item, usize>
188	0	where
189	0	I: Iterator,
190	0	I::Item: Clone,
191		{
192	0	let done = if self.first {
193	0	self.init()
194		} else {
195	0	self.increment_indices()
196		};
197	0	if done {
198	0	return Err(0);
199	0	}
200	0	for i in 0..n {
201	0	if self.increment_indices() {
202	0	return Err(i + 1);
203	0	}
204		}
205	0	Ok(self.indices.extract_item(&self.pool))
206	0	}
207		}
208
209		impl<I, Idx> Iterator for CombinationsGeneric<I, Idx>
210		where
211		I: Iterator,
212		I::Item: Clone,
213		Idx: PoolIndex<I::Item>,
214		{
215		type Item = Idx::Item;
216	0	fn next(&mut self) -> Option<Self::Item> {
217	0	let done = if self.first {
218	0	self.init()
219		} else {
220	0	self.increment_indices()
221		};
222
223	0	if done {
224	0	return None;
225	0	}
226
227	0	Some(self.indices.extract_item(&self.pool))
228	0	}
229
230	0	fn nth(&mut self, n: usize) -> Option<Self::Item> {
231	0	self.try_nth(n).ok()
232	0	}
233
234	0	fn size_hint(&self) -> (usize, Option<usize>) {
235	0	let (mut low, mut upp) = self.pool.size_hint();
236	0	low = remaining_for(low, self.first, self.indices.borrow()).unwrap_or(usize::MAX);
237	0	upp = upp.and_then(\|upp\| remaining_for(upp, self.first, self.indices.borrow()));
238	0	(low, upp)
239	0	}
240
241		#[inline]
242	0	fn count(self) -> usize {
243	0	self.n_and_count().1
244	0	}
245		}
246
247		impl<I, Idx> FusedIterator for CombinationsGeneric<I, Idx>
248		where
249		I: Iterator,
250		I::Item: Clone,
251		Idx: PoolIndex<I::Item>,
252		{
253		}
254
255		impl<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	0	pub(crate) fn reset(&mut self, k: usize) {
261	0	self.first = true;
262
263	0	if k < self.indices.len() {
264	0	self.indices.truncate(k);
265	0	for i in 0..k {
266	0	self.indices[i] = i;
267	0	}
268		} else {
269	0	for i in 0..self.indices.len() {
270	0	self.indices[i] = i;
271	0	}
272	0	self.indices.extend(self.indices.len()..k);
273	0	self.pool.prefill(k);
274		}
275	0	}
276		}
277
278		/// For a given size `n`, return the count of remaining combinations or None if it would overflow.
279	0	fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
280	0	let k = indices.len();
281	0	if n < k {
282	0	Some(0)
283	0	} else if first {
284	0	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	0	indices.iter().enumerate().try_fold(0usize, \|sum, (i, n0)\| {
305	0	sum.checked_add(checked_binomial(n - 1 - *n0, k - i)?)
306	0	})
307		}
308	0	}