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