itertools/
combinations.rs

1use std::fmt;
2use std::iter::FusedIterator;
3
4use super::lazy_buffer::LazyBuffer;
5use 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"]
11pub struct Combinations<I: Iterator> {
12    indices: Vec<usize>,
13    pool: LazyBuffer<I>,
14    first: bool,
15}
16
17impl<I> Clone for Combinations<I>
18    where I: Clone + Iterator,
19          I::Item: Clone,
20{
21    clone_fields!(indices, pool, first);
22}
23
24impl<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.
32pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
33    where I: Iterator
34{
35    let mut pool = LazyBuffer::new(iter);
36    pool.prefill(k);
37
38    Combinations {
39        indices: (0..k).collect(),
40        pool,
41        first: true,
42    }
43}
44
45impl<I: Iterator> Combinations<I> {
46    /// Returns the length of a combination produced by this iterator.
47    #[inline]
48    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    pub fn n(&self) -> usize { self.pool.len() }
54
55    /// Returns a reference to the source iterator.
56    #[inline]
57    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    pub(crate) fn reset(&mut self, k: usize) {
64        self.first = true;
65
66        if k < self.indices.len() {
67            self.indices.truncate(k);
68            for i in 0..k {
69                self.indices[i] = i;
70            }
71
72        } else {
73            for i in 0..self.indices.len() {
74                self.indices[i] = i;
75            }
76            self.indices.extend(self.indices.len()..k);
77            self.pool.prefill(k);
78        }
79    }
80}
81
82impl<I> Iterator for Combinations<I>
83    where I: Iterator,
84          I::Item: Clone
85{
86    type Item = Vec<I::Item>;
87    fn next(&mut self) -> Option<Self::Item> {
88        if self.first {
89            if self.k() > self.n() {
90                return None;
91            }
92            self.first = false;
93        } else if self.indices.is_empty() {
94            return None;
95        } else {
96            // Scan from the end, looking for an index to increment
97            let mut i: usize = self.indices.len() - 1;
98
99            // Check if we need to consume more from the iterator
100            if self.indices[i] == self.pool.len() - 1 {
101                self.pool.get_next(); // may change pool size
102            }
103
104            while self.indices[i] == i + self.pool.len() - self.indices.len() {
105                if i > 0 {
106                    i -= 1;
107                } else {
108                    // Reached the last combination
109                    return None;
110                }
111            }
112
113            // Increment index, and reset the ones to its right
114            self.indices[i] += 1;
115            for j in i+1..self.indices.len() {
116                self.indices[j] = self.indices[j - 1] + 1;
117            }
118        }
119
120        // Create result vector based on the indices
121        Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect())
122    }
123}
124
125impl<I> FusedIterator for Combinations<I>
126    where I: Iterator,
127          I::Item: Clone
128{}
itertools/combinations.rs

itertools/
combinations.rs
