itertools/
combinations.rs

Help
1
2use std::ops::Index;
3use std::fmt;
4
5/// An iterator to iterate through all the `n`-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"]
9pub struct Combinations<I: Iterator> {
10    n: usize,
11    indices: Vec<usize>,
12    pool: LazyBuffer<I>,
13    first: bool,
14}
15
16impl<I> fmt::Debug for Combinations<I>
17    where I: Iterator + fmt::Debug,
18          I::Item: fmt::Debug,
19{
20    debug_fmt_fields!(Combinations, n, indices, pool, first);
21}
22
23/// Create a new `Combinations` from a clonable iterator.
24pub fn combinations<I>(iter: I, n: usize) -> Combinations<I>
25    where I: Iterator
26{
27    let mut indices: Vec<usize> = Vec::with_capacity(n);
28    for i in 0..n {
29        indices.push(i);
30    }
31    let mut pool: LazyBuffer<I> = LazyBuffer::new(iter);
32
33    for _ in 0..n {
34        if !pool.get_next() {
35            break;
36        }
37    }
38
39    Combinations {
40        n: n,
41        indices: indices,
42        pool: pool,
43        first: true,
44    }
45}
46
47impl<I> Iterator for Combinations<I>
48    where I: Iterator,
49          I::Item: Clone
50{
51    type Item = Vec<I::Item>;
52    fn next(&mut self) -> Option<Self::Item> {
53        let mut pool_len = self.pool.len();
54        if self.pool.is_done() {
55            if pool_len == 0 || self.n > pool_len {
56                return None;
57            }
58        }
59
60        if self.first {
61            self.first = false;
62        } else if self.n == 0 {
63            return None;
64        } else {
65            // Scan from the end, looking for an index to increment
66            let mut i: usize = self.n - 1;
67
68            // Check if we need to consume more from the iterator
69            if self.indices[i] == pool_len - 1 && !self.pool.is_done() {
70                if self.pool.get_next() {
71                    pool_len += 1;
72                }
73            }
74
75            while self.indices[i] == i + pool_len - self.n {
76                if i > 0 {
77                    i -= 1;
78                } else {
79                    // Reached the last combination
80                    return None;
81                }
82            }
83
84            // Increment index, and reset the ones to its right
85            self.indices[i] += 1;
86            let mut j = i + 1;
87            while j < self.n {
88                self.indices[j] = self.indices[j - 1] + 1;
89                j += 1;
90            }
91        }
92
93        // Create result vector based on the indices
94        let mut result = Vec::with_capacity(self.n);
95        for i in self.indices.iter() {
96            result.push(self.pool[*i].clone());
97        }
98        Some(result)
99    }
100}
101
102#[derive(Debug)]
103struct LazyBuffer<I: Iterator> {
104    it: I,
105    done: bool,
106    buffer: Vec<I::Item>,
107}
108
109impl<I> LazyBuffer<I>
110    where I: Iterator
111{
112    pub fn new(it: I) -> LazyBuffer<I> {
113        let mut it = it;
114        let mut buffer = Vec::new();
115        let done;
116        if let Some(first) = it.next() {
117            buffer.push(first);
118            done = false;
119        } else {
120            done = true;
121        }
122        LazyBuffer {
123            it: it,
124            done: done,
125            buffer: buffer,
126        }
127    }
128
129    pub fn len(&self) -> usize {
130        self.buffer.len()
131    }
132
133    pub fn is_done(&self) -> bool {
134        self.done
135    }
136
137    pub fn get_next(&mut self) -> bool {
138        if self.done {
139            return false;
140        }
141        let next_item = self.it.next();
142        match next_item {
143            Some(x) => {
144                self.buffer.push(x);
145                true
146            }
147            None => {
148                self.done = true;
149                false
150            }
151        }
152    }
153}
154
155impl<I> Index<usize> for LazyBuffer<I>
156    where I: Iterator,
157          I::Item: Sized
158{
159    type Output = I::Item;
160
161    fn index<'b>(&'b self, _index: usize) -> &'b I::Item {
162        self.buffer.index(_index)
163    }
164}
165
itertools/combinations.rs

itertools/
combinations.rs
