1 use std::fmt;
2 use std::iter::FusedIterator;
3
4 use super::lazy_buffer::LazyBuffer;
5 use alloc::vec::Vec;
6
7 /// An iterator to iterate through all the `k`-length combinations in an iterator.
8 ///
9 /// See [`.combinations()`](crate::Itertools::combinations) for more information.
10 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
11 pub struct Combinations<I: Iterator> {
12 indices: Vec<usize>,
13 pool: LazyBuffer<I>,
14 first: bool,
15 }
16
17 impl<I> Clone for Combinations<I>
18 where I: Clone + Iterator,
19 I::Item: Clone,
20 {
21 clone_fields!(indices, pool, first);
22 }
23
24 impl<I> fmt::Debug for Combinations<I>
25 where I: Iterator + fmt::Debug,
26 I::Item: fmt::Debug,
27 {
28 debug_fmt_fields!(Combinations, indices, pool, first);
29 }
30
31 /// Create a new `Combinations` from a clonable iterator.
combinations<I>(iter: I, k: usize) -> Combinations<I> where I: Iterator32 pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
33 where I: Iterator
34 {
35 let mut pool = LazyBuffer::new(iter);
36 pool.prefill(k);
37
38 Combinations {
39 indices: (0..k).collect(),
40 pool,
41 first: true,
42 }
43 }
44
45 impl<I: Iterator> Combinations<I> {
46 /// Returns the length of a combination produced by this iterator.
47 #[inline]
k(&self) -> usize48 pub fn k(&self) -> usize { self.indices.len() }
49
50 /// Returns the (current) length of the pool from which combination elements are
51 /// selected. This value can change between invocations of [`next`](Combinations::next).
52 #[inline]
n(&self) -> usize53 pub fn n(&self) -> usize { self.pool.len() }
54
55 /// Returns a reference to the source iterator.
56 #[inline]
src(&self) -> &I57 pub(crate) fn src(&self) -> &I { &self.pool.it }
58
59 /// Resets this `Combinations` back to an initial state for combinations of length
60 /// `k` over the same pool data source. If `k` is larger than the current length
61 /// of the data pool an attempt is made to prefill the pool so that it holds `k`
62 /// elements.
reset(&mut self, k: usize)63 pub(crate) fn reset(&mut self, k: usize) {
64 self.first = true;
65
66 if k < self.indices.len() {
67 self.indices.truncate(k);
68 for i in 0..k {
69 self.indices[i] = i;
70 }
71
72 } else {
73 for i in 0..self.indices.len() {
74 self.indices[i] = i;
75 }
76 self.indices.extend(self.indices.len()..k);
77 self.pool.prefill(k);
78 }
79 }
80 }
81
82 impl<I> Iterator for Combinations<I>
83 where I: Iterator,
84 I::Item: Clone
85 {
86 type Item = Vec<I::Item>;
next(&mut self) -> Option<Self::Item>87 fn next(&mut self) -> Option<Self::Item> {
88 if self.first {
89 if self.k() > self.n() {
90 return None;
91 }
92 self.first = false;
93 } else if self.indices.is_empty() {
94 return None;
95 } else {
96 // Scan from the end, looking for an index to increment
97 let mut i: usize = self.indices.len() - 1;
98
99 // Check if we need to consume more from the iterator
100 if self.indices[i] == self.pool.len() - 1 {
101 self.pool.get_next(); // may change pool size
102 }
103
104 while self.indices[i] == i + self.pool.len() - self.indices.len() {
105 if i > 0 {
106 i -= 1;
107 } else {
108 // Reached the last combination
109 return None;
110 }
111 }
112
113 // Increment index, and reset the ones to its right
114 self.indices[i] += 1;
115 for j in i+1..self.indices.len() {
116 self.indices[j] = self.indices[j - 1] + 1;
117 }
118 }
119
120 // Create result vector based on the indices
121 Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect())
122 }
123 }
124
125 impl<I> FusedIterator for Combinations<I>
126 where I: Iterator,
127 I::Item: Clone
128 {}
129