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; |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 3 | use std::iter::FusedIterator; |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 4 | |
| 5 | use super::lazy_buffer::LazyBuffer; |
| 6 | |
| 7 | /// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement. |
| 8 | /// |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 9 | /// See [`.combinations_with_replacement()`](crate::Itertools::combinations_with_replacement) |
| 10 | /// for more information. |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 11 | #[derive(Clone)] |
| 12 | pub struct CombinationsWithReplacement<I> |
| 13 | where |
| 14 | I: Iterator, |
| 15 | I::Item: Clone, |
| 16 | { |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 17 | indices: Vec<usize>, |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 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 | { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 27 | debug_fmt_fields!(Combinations, indices, pool, first); |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 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 | { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 47 | let indices: Vec<usize> = alloc::vec![0; k]; |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 48 | let pool: LazyBuffer<I> = LazyBuffer::new(iter); |
| 49 | |
| 50 | CombinationsWithReplacement { |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 51 | indices, |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 52 | pool, |
| 53 | first: true, |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | impl<I> Iterator for CombinationsWithReplacement<I> |
| 58 | where |
| 59 | I: Iterator, |
| 60 | I::Item: Clone, |
| 61 | { |
| 62 | type Item = Vec<I::Item>; |
| 63 | fn next(&mut self) -> Option<Self::Item> { |
| 64 | // If this is the first iteration, return early |
| 65 | if self.first { |
| 66 | // In empty edge cases, stop iterating immediately |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 67 | return if self.indices.len() != 0 && !self.pool.get_next() { |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 68 | None |
| 69 | // Otherwise, yield the initial state |
| 70 | } else { |
| 71 | self.first = false; |
| 72 | Some(self.current()) |
| 73 | }; |
| 74 | } |
| 75 | |
| 76 | // Check if we need to consume more from the iterator |
| 77 | // This will run while we increment our first index digit |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 78 | self.pool.get_next(); |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 79 | |
| 80 | // Work out where we need to update our indices |
| 81 | let mut increment: Option<(usize, usize)> = None; |
| 82 | for (i, indices_int) in self.indices.iter().enumerate().rev() { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 83 | if *indices_int < self.pool.len()-1 { |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 84 | increment = Some((i, indices_int + 1)); |
| 85 | break; |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | match increment { |
| 90 | // If we can update the indices further |
| 91 | Some((increment_from, increment_value)) => { |
| 92 | // We need to update the rightmost non-max value |
| 93 | // and all those to the right |
| 94 | for indices_index in increment_from..self.indices.len() { |
| 95 | self.indices[indices_index] = increment_value |
| 96 | } |
| 97 | Some(self.current()) |
| 98 | } |
| 99 | // Otherwise, we're done |
| 100 | None => None, |
| 101 | } |
| 102 | } |
| 103 | } |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 104 | |
| 105 | impl<I> FusedIterator for CombinationsWithReplacement<I> |
| 106 | where |
| 107 | I: Iterator, |
| 108 | I::Item: Clone, |
| 109 | {} |