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