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
99impl<I> Iterator for Combinations<I>
100where
101 I: Iterator,
102 I::Item: Clone,
103{
104 type Item = Vec<I::Item>;
105 fn next(&mut self) -> Option<Self::Item> {
106 if self.first {
107 self.pool.prefill(self.k());
108 if self.k() > self.n() {
109 return None;
110 }
111 self.first = false;
112 } else if self.indices.is_empty() {
113 return None;
114 } else {
115 // Scan from the end, looking for an index to increment
116 let mut i: usize = self.indices.len() - 1;
117
118 // Check if we need to consume more from the iterator
119 if self.indices[i] == self.pool.len() - 1 {
120 self.pool.get_next(); // may change pool size
121 }
122
123 while self.indices[i] == i + self.pool.len() - self.indices.len() {
124 if i > 0 {
125 i -= 1;
126 } else {
127 // Reached the last combination
128 return None;
129 }
130 }
131
132 // Increment index, and reset the ones to its right
133 self.indices[i] += 1;
134 for j in i + 1..self.indices.len() {
135 self.indices[j] = self.indices[j - 1] + 1;
136 }
137 }
138
139 // Create result vector based on the indices
140 Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect())
141 }
142
143 fn size_hint(&self) -> (usize, Option<usize>) {
144 let (mut low, mut upp) = self.pool.size_hint();
145 low = remaining_for(low, self.first, &self.indices).unwrap_or(usize::MAX);
146 upp = upp.and_then(|upp| remaining_for(upp, self.first, &self.indices));
147 (low, upp)
148 }
149
150 #[inline]
151 fn count(self) -> usize {
152 self.n_and_count().1
153 }
154}
155
156impl<I> FusedIterator for Combinations<I>
157where
158 I: Iterator,
159 I::Item: Clone,
160{
161}
162
163/// For a given size `n`, return the count of remaining combinations or None if it would overflow.
164fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
165 let k = indices.len();
166 if n < k {
167 Some(0)
168 } else if first {
169 checked_binomial(n, k)
170 } else {
171 // https://en.wikipedia.org/wiki/Combinatorial_number_system
172 // http://www.site.uottawa.ca/~lucia/courses/5165-09/GenCombObj.pdf
173
174 // The combinations generated after the current one can be counted by counting as follows:
175 // - The subsequent combinations that differ in indices[0]:
176 // If subsequent combinations differ in indices[0], then their value for indices[0]
177 // must be at least 1 greater than the current indices[0].
178 // As indices is strictly monotonically sorted, this means we can effectively choose k values
179 // from (n - 1 - indices[0]), leading to binomial(n - 1 - indices[0], k) possibilities.
180 // - The subsequent combinations with same indices[0], but differing indices[1]:
181 // Here we can choose k - 1 values from (n - 1 - indices[1]) values,
182 // leading to binomial(n - 1 - indices[1], k - 1) possibilities.
183 // - (...)
184 // - The subsequent combinations with same indices[0..=i], but differing indices[i]:
185 // Here we can choose k - i values from (n - 1 - indices[i]) values: binomial(n - 1 - indices[i], k - i).
186 // Since subsequent combinations can in any index, we must sum up the aforementioned binomial coefficients.
187
188 // Below, `n0` resembles indices[i].
189 indices.iter().enumerate().try_fold(0usize, |sum, (i, n0)| {
190 sum.checked_add(checked_binomial(n - 1 - *n0, k - i)?)
191 })
192 }
193}
194