Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 1 | use std::fmt; |
| 2 | |
| 3 | use super::lazy_buffer::LazyBuffer; |
| 4 | |
| 5 | /// An iterator to iterate through all the `k`-length combinations in an iterator. |
| 6 | /// |
| 7 | /// See [`.combinations()`](../trait.Itertools.html#method.combinations) for more information. |
| 8 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] |
| 9 | pub struct Combinations<I: Iterator> { |
| 10 | indices: Vec<usize>, |
| 11 | pool: LazyBuffer<I>, |
| 12 | first: bool, |
| 13 | } |
| 14 | |
| 15 | impl<I> Clone for Combinations<I> |
| 16 | where I: Clone + Iterator, |
| 17 | I::Item: Clone, |
| 18 | { |
| 19 | clone_fields!(indices, pool, first); |
| 20 | } |
| 21 | |
| 22 | impl<I> fmt::Debug for Combinations<I> |
| 23 | where I: Iterator + fmt::Debug, |
| 24 | I::Item: fmt::Debug, |
| 25 | { |
| 26 | debug_fmt_fields!(Combinations, indices, pool, first); |
| 27 | } |
| 28 | |
| 29 | /// Create a new `Combinations` from a clonable iterator. |
| 30 | pub fn combinations<I>(iter: I, k: usize) -> Combinations<I> |
| 31 | where I: Iterator |
| 32 | { |
| 33 | let mut pool: LazyBuffer<I> = LazyBuffer::new(iter); |
| 34 | |
| 35 | for _ in 0..k { |
| 36 | if !pool.get_next() { |
| 37 | break; |
| 38 | } |
| 39 | } |
| 40 | |
| 41 | Combinations { |
| 42 | indices: (0..k).collect(), |
| 43 | pool, |
| 44 | first: true, |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | impl<I> Iterator for Combinations<I> |
| 49 | where I: Iterator, |
| 50 | I::Item: Clone |
| 51 | { |
| 52 | type Item = Vec<I::Item>; |
| 53 | fn next(&mut self) -> Option<Self::Item> { |
| 54 | if self.first { |
| 55 | if self.pool.is_done() { |
| 56 | return None; |
| 57 | } |
| 58 | self.first = false; |
| 59 | } else if self.indices.len() == 0 { |
| 60 | return None; |
| 61 | } else { |
| 62 | // Scan from the end, looking for an index to increment |
| 63 | let mut i: usize = self.indices.len() - 1; |
| 64 | |
| 65 | // Check if we need to consume more from the iterator |
| 66 | if self.indices[i] == self.pool.len() - 1 { |
| 67 | self.pool.get_next(); // may change pool size |
| 68 | } |
| 69 | |
| 70 | while self.indices[i] == i + self.pool.len() - self.indices.len() { |
| 71 | if i > 0 { |
| 72 | i -= 1; |
| 73 | } else { |
| 74 | // Reached the last combination |
| 75 | return None; |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | // Increment index, and reset the ones to its right |
| 80 | self.indices[i] += 1; |
| 81 | for j in i+1..self.indices.len() { |
| 82 | self.indices[j] = self.indices[j - 1] + 1; |
| 83 | } |
| 84 | } |
| 85 | |
| 86 | // Create result vector based on the indices |
| 87 | Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect()) |
| 88 | } |
| 89 | } |