aoclib_rs/
iter.rs

1use std::iter;
2
3use crate::inc_selector;
4
5/// Forward/reverse inclusive range. If `start <= end`, returns a forward iterator from `start`
6/// to `end`, inclusive. Else, returns a reverse iterator from `end` to `start`, inclusive.
7///
8/// ```
9/// use aoclib_rs::iter::fwd_rev_incl_range;
10/// assert_eq!(fwd_rev_incl_range(1, 3).collect::<Vec<_>>(), [1, 2, 3]);
11/// assert_eq!(fwd_rev_incl_range(3, 1).collect::<Vec<_>>(), [3, 2, 1]);
12/// ```
13pub fn fwd_rev_incl_range(start: usize, end: usize) -> impl Iterator<Item = usize> {
14    let mut fwd = start..=end;
15    let mut rev = (end..=start).rev();
16
17    iter::from_fn(move || if start > end { rev.next() } else { fwd.next() })
18}
19
20/// Iterates through every pair in the input slice. Returns references to the input.
21///
22/// ```
23/// let input = [1, 2, 3];
24/// let want = [(&input[0], &input[1]), (&input[0], &input[2]), (&input[1], &input[2])];
25/// assert_eq!(aoclib_rs::iter::pairwise_iter(&input).collect::<Vec<_>>(), want);
26/// ```
27pub fn pairwise_iter<T>(v: &[T]) -> impl Iterator<Item = (&T, &T)> {
28    let mut it = pairwise_iter_i(v);
29    iter::from_fn(move || it.next().map(|(i, j)| (&v[i], &v[j])))
30}
31
32/// Iterates through every pair in the input slice. Returns copies of the input.
33///
34/// ```
35/// let input = [1, 2, 3];
36/// let want = [(1, 2), (1, 3), (2, 3)];
37/// assert_eq!(aoclib_rs::iter::pairwise_iter_copy(&input).collect::<Vec<_>>(), want);
38/// ```
39pub fn pairwise_iter_copy<T: Copy>(v: &[T]) -> impl Iterator<Item = (T, T)> {
40    let mut it = pairwise_iter_i(v);
41    iter::from_fn(move || it.next().map(|(i, j)| (v[i], v[j])))
42}
43
44/// Iterates through every pair in the input slice. Returns indices into the input.
45///
46/// ```
47/// let input = [1, 2, 3];
48/// let want = [(0, 1), (0, 2), (1, 2)];
49/// assert_eq!(aoclib_rs::iter::pairwise_iter_i(&input).collect::<Vec<_>>(), want);
50/// ```
51pub fn pairwise_iter_i<T>(v: &[T]) -> impl Iterator<Item = (usize, usize)> {
52    let mut i = 0;
53    let mut j = 0;
54    iter::from_fn(move || {
55        if v.len() < 2 {
56            return None;
57        }
58
59        j += 1;
60        if j >= v.len() {
61            i += 1;
62            j = i + 1;
63        }
64
65        if j >= v.len() {
66            return None;
67        }
68
69        Some((i, j))
70    })
71}
72
73/// Iterates through a "selector" of `len` elements, generating every possible selection. The idea
74/// is that you "select" the elements in some other structure at the indices where the returned `Vec`
75/// is true. Put another way, the returned `Vec<bool>`s can be looked at as counting up in binary.
76///
77/// ```
78/// assert_eq!(
79///     aoclib_rs::iter::selector_iter(2).collect::<Vec<_>>(),
80///     [
81///         vec![false, false],
82///         vec![false, true],
83///         vec![true, false],
84///         vec![true, true]
85///     ]
86/// );
87/// ```
88pub fn selector_iter(len: usize) -> impl Iterator<Item = Vec<bool>> {
89    let mut v: Vec<bool> = vec![false; len];
90    let mut first = true;
91    iter::from_fn(move || {
92        if first {
93            first = false;
94            return Some(v.clone());
95        }
96        if inc_selector(&mut v) {
97            None
98        } else {
99            Some(v.clone())
100        }
101    })
102}
103
104// TODO: does this make more sense as sets instead of Vecs?
105// TODO: do we need a copy version of this?
106/// Iterates through every possible combination of selecting elements from the input. Returns
107/// references to the input.
108///
109/// ```
110/// let input = [1, 2];
111/// let want = [vec![], vec![&input[1]], vec![&input[0]], vec![&input[0], &input[1]]];
112/// assert_eq!(aoclib_rs::iter::selection_iter(&input).collect::<Vec<_>>(), want);
113/// ```
114pub fn selection_iter<T>(v: &[T]) -> impl Iterator<Item = Vec<&T>> {
115    let mut selector: Vec<bool> = vec![false; v.len()];
116    let mut first = true;
117    iter::from_fn(move || {
118        if first {
119            first = false;
120            return Some(Vec::new());
121        }
122        if inc_selector(&mut selector) {
123            None
124        } else {
125            Some(
126                v.iter()
127                    .enumerate()
128                    .filter_map(|(i, e)| if selector[i] { Some(e) } else { None })
129                    .collect(),
130            )
131        }
132    })
133}
134
135// Non-recursive variant of Heap's Algorithm https://en.wikipedia.org/wiki/Heap%27s_algorithm
136/// Iterates through all permutations of the input. Returns copies of the input.
137///
138/// ```
139/// assert_eq!(
140///     aoclib_rs::iter::permutation_iter_copy(vec![1, 2, 3]).collect::<Vec<_>>(),
141///     [
142///         vec![1, 2, 3],
143///         vec![2, 1, 3],
144///         vec![3, 1, 2],
145///         vec![1, 3, 2],
146///         vec![2, 3, 1],
147///         vec![3, 2, 1]
148///     ]
149/// );
150/// ```
151pub fn permutation_iter_copy<T: Copy>(mut a: Vec<T>) -> impl Iterator<Item = Vec<T>> {
152    let n = a.len();
153    let mut c: Vec<usize> = vec![0; n];
154    let mut first = true;
155    let mut i = 1;
156    iter::from_fn(move || {
157        if first {
158            first = false;
159            return Some(a.clone());
160        }
161
162        while i < n {
163            if c[i] < i {
164                if i % 2 == 0 {
165                    (a[0], a[i]) = (a[i], a[0]);
166                } else {
167                    (a[c[i]], a[i]) = (a[i], a[c[i]]);
168                }
169                c[i] += 1;
170                i = 1;
171                return Some(a.clone());
172            } else {
173                c[i] = 0;
174                i += 1;
175            }
176        }
177
178        None
179    })
180}