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