]>
Commit | Line | Data |
---|---|---|
f9f354fc XL |
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 | |
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. | |
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>; | |
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 | } |