1 use std::fmt;
2
3 use super::lazy_buffer::LazyBuffer;
4
5 /// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement.
6 ///
7 /// See [`.combinations_with_replacement()`](../trait.Itertools.html#method.combinations_with_replacement) for more information.
8 #[derive(Clone)]
9 pub struct CombinationsWithReplacement<I>
10 where
11 I: Iterator,
12 I::Item: Clone,
13 {
14 k: usize,
15 indices: Vec<usize>,
16 // The current known max index value. This increases as pool grows.
17 max_index: usize,
18 pool: LazyBuffer<I>,
19 first: bool,
20 }
21
22 impl<I> fmt::Debug for CombinationsWithReplacement<I>
23 where
24 I: Iterator + fmt::Debug,
25 I::Item: fmt::Debug + Clone,
26 {
27 debug_fmt_fields!(Combinations, k, indices, max_index, pool, first);
28 }
29
30 impl<I> CombinationsWithReplacement<I>
31 where
32 I: Iterator,
33 I::Item: Clone,
34 {
35 /// Map the current mask over the pool to get an output combination
current(&self) -> Vec<I::Item>36 fn current(&self) -> Vec<I::Item> {
37 self.indices.iter().map(|i| self.pool[*i].clone()).collect()
38 }
39 }
40
41 /// Create a new `CombinationsWithReplacement` from a clonable iterator.
combinations_with_replacement<I>(iter: I, k: usize) -> CombinationsWithReplacement<I> where I: Iterator, I::Item: Clone,42 pub fn combinations_with_replacement<I>(iter: I, k: usize) -> CombinationsWithReplacement<I>
43 where
44 I: Iterator,
45 I::Item: Clone,
46 {
47 let indices: Vec<usize> = vec![0; k];
48 let pool: LazyBuffer<I> = LazyBuffer::new(iter);
49
50 CombinationsWithReplacement {
51 k,
52 indices,
53 max_index: 0,
54 pool,
55 first: true,
56 }
57 }
58
59 impl<I> Iterator for CombinationsWithReplacement<I>
60 where
61 I: Iterator,
62 I::Item: Clone,
63 {
64 type Item = Vec<I::Item>;
next(&mut self) -> Option<Self::Item>65 fn next(&mut self) -> Option<Self::Item> {
66 // If this is the first iteration, return early
67 if self.first {
68 // In empty edge cases, stop iterating immediately
69 return if self.k != 0 && !self.pool.get_next() {
70 None
71 // Otherwise, yield the initial state
72 } else {
73 self.first = false;
74 Some(self.current())
75 };
76 }
77
78 // Check if we need to consume more from the iterator
79 // This will run while we increment our first index digit
80 if self.pool.get_next() {
81 self.max_index = self.pool.len() - 1;
82 }
83
84 // Work out where we need to update our indices
85 let mut increment: Option<(usize, usize)> = None;
86 for (i, indices_int) in self.indices.iter().enumerate().rev() {
87 if indices_int < &self.max_index {
88 increment = Some((i, indices_int + 1));
89 break;
90 }
91 }
92
93 match increment {
94 // If we can update the indices further
95 Some((increment_from, increment_value)) => {
96 // We need to update the rightmost non-max value
97 // and all those to the right
98 for indices_index in increment_from..self.indices.len() {
99 self.indices[indices_index] = increment_value
100 }
101 Some(self.current())
102 }
103 // Otherwise, we're done
104 None => None,
105 }
106 }
107 }
108