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}