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, indices:array::from_fn(|i: usize| 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(len:k);
265 for i: usize in 0..k {
266 self.indices[i] = i;
267 }
268 } else {
269 for i: usize in 0..self.indices.len() {
270 self.indices[i] = i;
271 }
272 self.indices.extend(self.indices.len()..k);
273 self.pool.prefill(len: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}
309