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