feanor_math/
iters.rs

1use std::{cmp::min, convert::TryInto, iter::FusedIterator};
2
3///
4/// Iterator over all combinations of the given set of the given size.
5/// See also [`combinations()`].
6/// 
7#[derive(Debug, Clone)]
8pub struct IterCombinations<I, F, T> 
9    where I: Iterator + Clone, I::Item: Clone, F: FnMut(&[I::Item]) -> T
10{
11    base: std::iter::Peekable<std::iter::Enumerate<I>>,
12    iterators: Box<[std::iter::Peekable<std::iter::Enumerate<I>>]>,
13    done: bool,
14    converter: F,
15    buffer: Box<[I::Item]>
16}
17
18impl<I, F, T> Iterator for IterCombinations<I, F, T> 
19    where I: Iterator + Clone, I::Item: Clone, F: FnMut(&[I::Item]) -> T
20{
21    type Item = T;
22
23    fn next(&mut self) -> Option<Self::Item> {
24        if self.done {
25            return None;
26        } else if self.iterators.len() == 0 {
27            self.done = true;
28            return Some((self.converter)(&[]));
29        }
30        let result = (self.converter)(&self.buffer[..]);
31        for i in 0..self.iterators.len() - 1 {
32            let (its, next_its) = &mut self.iterators[i..i+2].split_at_mut(1);
33            let (it, next_it) = (&mut its[0], &mut next_its[0]);
34            let can_forward = {
35                let (index, _) = it.peek().unwrap();
36                let (next_index, _) = next_it.peek().unwrap();
37                index + 1 < *next_index
38            };
39            if can_forward {
40                _ = it.next().unwrap();
41                self.buffer[i] = it.peek().unwrap().1.clone();
42                return Some(result);
43            } else {
44                // reset and continue with next iterator
45                *it = self.base.clone();
46                for _ in 0..i {
47                    _ = it.next();
48                }
49                let (_, x) = it.peek().unwrap();
50                self.buffer[i] = x.clone();
51            }
52        }
53        if let Some(last_it) = self.iterators.last_mut() {
54            _ = last_it.next();
55            if let Some(x) = last_it.peek() {
56                *self.buffer.last_mut().unwrap() = x.1.clone();
57            } else {
58                self.done = true;
59            }
60        }
61        return Some(result);
62    }
63}
64
65impl<I, F, T> FusedIterator for IterCombinations<I, F, T> 
66    where I: Iterator + Clone, I::Item: Clone, F: FnMut(&[I::Item]) -> T {}
67
68///
69/// Returns an iterator over all combinations of the elements yielded by the base iterator of size `k`.
70/// 
71/// Since the items yielded by the iterator are collections of dynamic length,
72/// creating them might be quite costly. In many applications, this is not really
73/// required, and so this function accepts a second argument - a converter function -
74/// that is called on each tuple in the cartesian product and its return value
75/// is yielded by the product iterator.
76/// 
77/// Note that clones of the given base iterator must all have the same iteration order.
78/// 
79/// # Example
80/// 
81/// ```rust
82/// # use feanor_math::iters::combinations;
83/// assert_eq!(
84///     vec![(1, 2), (1, 3), (2, 3)],
85///     combinations(
86///         [1, 2, 3].into_iter(), 
87///         2, 
88///         |a| (a[0], a[1])
89///     ).collect::<Vec<_>>()
90/// );
91/// ```
92/// 
93pub fn combinations<I, F, T>(it: I, k: usize, converter: F) -> IterCombinations<I, F, T> 
94    where I: Iterator + Clone, I::Item: Clone, F: FnMut(&[I::Item]) -> T
95{
96    let enumerated_it = it.enumerate().peekable();
97    let mut start_iterators = Vec::with_capacity(k);
98    let mut buffer = Vec::with_capacity(k);
99    let mut start_it = enumerated_it.clone();
100    for _ in 0..k {
101        start_iterators.push(start_it.clone());
102        if start_it.peek().is_none() {
103            return IterCombinations {
104                base: enumerated_it,
105                iterators: start_iterators.into_boxed_slice(),
106                done: true,
107                converter: converter,
108                buffer: buffer.into_boxed_slice()
109            };
110        }
111        let (_, x) = start_it.next().unwrap();
112        buffer.push(x);
113    }
114    return IterCombinations {
115        base: enumerated_it,
116        iterators: start_iterators.into_boxed_slice(),
117        done: false,
118        converter: converter,
119        buffer: buffer.into_boxed_slice()
120    };
121}
122
123///
124/// Clones every element of the given slice and returns an owned pointer to the copy.
125/// 
126pub fn clone_slice<T>(slice: &[T]) -> Box<[T]> 
127    where T: Clone
128{
129    let vec: Vec<T> = slice.iter().cloned().collect();
130    return vec.into_boxed_slice();
131}
132
133///
134/// Clones every element of the given array.
135/// 
136pub fn clone_array<T, const N: usize>(slice: &[T]) -> [T; N] 
137    where T: Copy
138{
139    slice.try_into().unwrap()
140}
141
142///
143/// Returns an iterator over all size-`k` subsets of the
144/// elements returned by the given iterator.
145/// 
146/// This is a special version of [`combinations()`], which
147/// accepts a converter function and thus can avoid cloning
148/// each subset.
149/// 
150pub fn basic_combinations<I>(it: I, k: usize) -> impl Iterator<Item = Box<[I::Item]>>
151    where I: Iterator + Clone, I::Item: Clone, 
152{
153    combinations(it, k, clone_slice)
154}
155
156///
157/// Calls the given `converter` function once for each subset of the elements
158/// returned by the given iterator, and returns an iterator over the
159/// outputs of the `converter` function.
160/// 
161pub fn powerset<I, F, T>(it: I, converter: F) -> impl Iterator<Item = T>
162    where I: Iterator + Clone, I::Item: Clone, F: Clone + FnMut(&[I::Item]) -> T
163{
164    let len = it.clone().count();
165    (0..=len).flat_map(move |i| combinations(it.clone(), i, converter.clone()))
166}
167
168///
169/// Returns an iterator over all subsets of the
170/// elements returned by the given iterator.
171/// 
172/// This is a special version of [`powerset()`], which
173/// accepts a converter function and thus can avoid cloning
174/// each subset.
175/// 
176pub fn basic_powerset<I>(it: I) -> impl Iterator<Item = Box<[I::Item]>>
177    where I: Iterator + Clone, I::Item: Clone
178{
179    powerset(it, clone_slice)
180}
181
182///
183/// Iterator over the multiset combinations of a set and a given size.
184/// See also [`multiset_combinations()`].
185/// 
186#[derive(Debug, Clone)]
187pub struct MultisetCombinations<'a, F, T>
188    where F: FnMut(&[usize]) -> T
189{
190    converter: F,
191    superset: &'a [usize],
192    current: Option<Box<[usize]>>,
193    last_moved: usize,
194    first_trailing_empty: usize
195}
196
197impl<'a, F, T> Iterator for MultisetCombinations<'a, F, T>
198    where F: FnMut(&[usize]) -> T
199{
200    type Item = T;
201
202    fn next(&mut self) -> Option<T> {
203        if self.current.is_none() {
204            return None;
205        }
206        let current = &mut self.current.as_mut().unwrap();
207        let result = (self.converter)(current);
208        let mut removed = 0;
209        let mut found_empty_place = self.last_moved + 1 < self.first_trailing_empty;
210        while !found_empty_place || current[self.last_moved] == 0 {
211            found_empty_place |= current[self.last_moved] < self.superset[self.last_moved];
212            removed += current[self.last_moved];
213            current[self.last_moved] = 0;
214            if self.last_moved == 0 {
215                self.current = None;
216                return Some(result);
217            }
218            self.last_moved -= 1;
219        }
220        removed += 1;
221        current[self.last_moved] -= 1;
222        while removed > 0 {
223            self.last_moved += 1;
224            if current[self.last_moved] + removed > self.superset[self.last_moved] {
225                removed = current[self.last_moved] + removed - self.superset[self.last_moved];
226                current[self.last_moved] = self.superset[self.last_moved];
227            } else {
228                current[self.last_moved] += removed;
229                removed = 0;
230            }
231        }
232        return Some(result);
233    }
234}
235
236impl<'a, F, T> FusedIterator for MultisetCombinations<'a, F, T>
237    where F: FnMut(&[usize]) -> T {}
238
239///
240/// Returns an iterator over all multi-subsets of a given set with a specified size.
241/// 
242/// Since the items yielded by the iterator are collections of dynamic length,
243/// creating them might be quite costly. In many applications, this is not really
244/// required, and so this function accepts a second argument - a converter function -
245/// that is called on each tuple in the cartesian product and its return value
246/// is yielded by the product iterator.
247/// 
248/// Note that clones of the given base iterator must all have the same iteration order.
249/// 
250/// This iterator returns its elements in descending lexicographic order, so the first
251/// yielded element is the largest one w.r.t. lexicographic order.
252/// 
253/// # Example
254/// ```rust
255/// # use feanor_math::iters::multiset_combinations;
256/// assert_eq!(
257///     vec![(1, 1, 0), (1, 0, 1), (0, 2, 0), (0, 1, 1), (0, 0, 2)],
258///     multiset_combinations(
259///         &[1, 2, 3], 
260///         2, 
261///         |a| (a[0], a[1], a[2])
262///     ).collect::<Vec<_>>()
263/// );
264/// ```
265/// ```rust
266/// # use feanor_math::iters::multiset_combinations;
267/// assert_eq!(
268///     vec![(2, 0, 0), (1, 1, 0), (1, 0, 1), (0, 2, 0), (0, 1, 1), (0, 0, 2)],
269///     multiset_combinations(
270///         &[2, 2, 2], 
271///         2, 
272///         |a| (a[0], a[1], a[2])
273///     ).collect::<Vec<_>>()
274/// );
275/// ```
276/// 
277pub fn multiset_combinations<'a, F, T>(multiset: &'a [usize], size: usize, converter: F) -> MultisetCombinations<'a, F, T>
278    where F: FnMut(&[usize]) -> T
279{
280    assert!(multiset.len() > 0);
281    if size > multiset.iter().copied().sum::<usize>() {
282        return MultisetCombinations {
283            converter: converter,
284            superset: multiset,
285            first_trailing_empty: 0,
286            current: None,
287            last_moved: 0
288        };
289    }
290    let mut start = (0..multiset.len()).map(|_| 0).collect::<Vec<_>>().into_boxed_slice();
291    let mut to_insert = size;
292    let mut last_moved = 0;
293    let mut i = 0;
294    while to_insert > 0 {
295        start[i] = min(multiset[i], to_insert);
296        last_moved = i;
297        to_insert -= start[i];
298        i += 1;
299    }
300    let mut trailing_empty_entries = 0;
301    while trailing_empty_entries < multiset.len() && multiset[multiset.len() - trailing_empty_entries - 1] == 0 {
302        trailing_empty_entries += 1;
303    }
304    return MultisetCombinations {
305        converter: converter,
306        superset: multiset,
307        first_trailing_empty: multiset.len() - trailing_empty_entries,
308        current: Some(start),
309        last_moved: last_moved
310    };
311}
312
313///
314/// Iterator over the elements of the cartesian product of multiple iterators.
315/// See also [`multi_cartesian_product()`].
316/// 
317#[derive(Debug)]
318pub struct MultiProduct<I, F, G, T> 
319    where I: Iterator + Clone, 
320        F: FnMut(&[I::Item]) -> T,
321        G: Clone + Fn(usize, &I::Item) -> I::Item
322{
323    base_iters: Vec<I>,
324    current_iters: Vec<I>,
325    current: Vec<I::Item>,
326    clone_el: G,
327    converter: F,
328    done: bool
329}
330
331impl<I, F, G, T> Clone for MultiProduct<I, F, G, T>
332    where I: Iterator + Clone, 
333        F: Clone + FnMut(&[I::Item]) -> T,
334        G: Clone + Fn(usize, &I::Item) -> I::Item
335{
336    fn clone(&self) -> Self {
337        MultiProduct {
338            base_iters: self.base_iters.clone(),
339            current_iters: self.current_iters.clone(),
340            current: self.current.iter().enumerate().map(|(i, x)| (self.clone_el)(i, x)).collect(),
341            clone_el: self.clone_el.clone(),
342            converter: self.converter.clone(),
343            done: self.done
344        }
345    }
346}
347
348impl<I, F, G, T> Iterator for MultiProduct<I, F, G, T>
349    where I: Iterator + Clone, 
350        F: FnMut(&[I::Item]) -> T,
351        G: Clone + Fn(usize, &I::Item) -> I::Item
352{
353    type Item = T;
354
355    fn next(&mut self) -> Option<Self::Item> {
356        if self.done {
357            return None;
358        }
359        let result = (self.converter)(&self.current[..]);
360        let mut i = self.base_iters.len();
361        self.done = true;
362        while i > 0 {
363            i = i - 1;
364            if let Some(val) = self.current_iters[i].next() {
365                self.current[i] = val;
366                self.done = false;
367                for j in (i + 1)..self.base_iters.len() {
368                    self.current_iters[j] = self.base_iters[j].clone();
369                    self.current[j] = self.current_iters[j].next().unwrap();
370                }
371                break;
372            }
373        }
374        return Some(result);
375    }
376}
377
378impl<I, F, G, T> FusedIterator for MultiProduct<I, F, G, T>
379    where I: Iterator + Clone, 
380        F: Clone + FnMut(&[I::Item]) -> T,
381        G: Clone + Fn(usize, &I::Item) -> I::Item
382{}
383
384///
385/// Creates an iterator that computes the cartesian product of the elements
386/// yielded by a number of iterators. These iterators are given by one iterator
387/// over them.
388/// 
389/// Since the items yielded by the iterator are collections of dynamic length,
390/// creating them might be quite costly. In many applications, this is not really
391/// required, and so this function accepts a second argument - a converter function -
392/// that is called on each tuple in the cartesian product and its return value
393/// is yielded by the product iterator.
394/// 
395/// Note that clones of the given base iterator must all have the same iteration order.
396/// 
397/// # Example
398/// ```rust
399/// # use feanor_math::iters::multi_cartesian_product;
400/// assert_eq!(
401///     vec![(2, 0), (2, 1), (2, 2), (3, 0), (3, 1), (3, 2)],
402///     multi_cartesian_product(
403///         vec![(2..=3), (0..=2)].into_iter(),
404///         |x| (x[0], x[1]),
405///         |_, x| *x
406///     ).collect::<Vec<_>>()
407/// );
408/// ```
409/// 
410pub fn multi_cartesian_product<J, F, G, T>(iters: J, converter: F, clone_el: G) -> MultiProduct<J::Item, F, G, T>
411    where J: Iterator, 
412        J::Item: Iterator + Clone, 
413        F: FnMut(&[<J::Item as Iterator>::Item]) -> T,
414        G: Clone + Fn(usize, &<J::Item as Iterator>::Item) -> <J::Item as Iterator>::Item
415{
416    let base_iters = iters.collect::<Vec<_>>();
417    let mut current_iters = base_iters.clone();
418    let mut current = Vec::with_capacity(current_iters.len());
419    for it in current_iters.iter_mut() {
420        if let Some(v) = it.next() {
421            current.push(v);
422        } else {
423            return MultiProduct {
424                done: true,
425                converter: converter,
426                base_iters: base_iters,
427                clone_el: clone_el,
428                current_iters: current_iters,
429                current: current
430            };
431        }
432    }
433    return MultiProduct {
434        done: false,
435        converter: converter,
436        base_iters: base_iters,
437        current_iters: current_iters,
438        clone_el: clone_el,
439        current: current
440    };
441}
442
443///
444/// An iterator that "condenses" a given iterator, i.e. combines a variable number
445/// of base iterator elements into a single element.
446/// 
447/// Use through the function [`condense()`].
448/// 
449#[derive(Debug)]
450pub struct CondenseIter<I, F, T>
451    where I: Iterator, F: FnMut(I::Item) -> Option<T>
452{
453    base_iter: I,
454    f: F
455}
456
457impl<I, F, T> Iterator for CondenseIter<I, F, T>
458    where I: Iterator, F: FnMut(I::Item) -> Option<T>
459{
460    type Item = T;
461
462    fn next(&mut self) -> Option<T> {
463        while let Some(el) = self.base_iter.next() {
464            if let Some(result) = (self.f)(el) {
465                return Some(result);
466            }
467        }
468        return None;
469    }
470}
471
472impl<I, F, T> FusedIterator for CondenseIter<I, F, T>
473    where I: FusedIterator, F: FnMut(I::Item) -> Option<T>
474{}
475
476///
477/// Creates an iterator that "condenses" the given iterator.
478/// Concretely, the new iterator processes a variable amount
479/// of input items to produce an output item.
480/// 
481/// This processing is done using the passed function. It is
482/// called repeatedly on input iterator items, until it returns
483/// some value, which is the next value yielded by the output
484/// iterator.
485/// 
486/// # Example
487/// ```rust
488/// # use feanor_math::iters::condense;
489/// let mut accumulator = 0;
490/// assert_eq!(vec![6, 9, 6, 7], condense(vec![1, 2, 3, 4, 5, 6, 7].into_iter(), move |a| {
491///     accumulator += a;
492///     if accumulator >= 5 {
493///         let result = accumulator;
494///         accumulator = 0;
495///         return Some(result);
496///     } else {
497///         return None;
498///     }
499/// }).collect::<Vec<i64>>());
500/// ```
501/// 
502pub fn condense<I, F, T>(iter: I, f: F) -> CondenseIter<I, F, T>
503    where I: Iterator, F: FnMut(I::Item) -> Option<T>
504{
505    CondenseIter { base_iter: iter, f: f }
506}
507
508#[test]
509fn test_converted_combinations() {
510    let a = [2, 3, 5, 7];
511    assert_eq!(1, combinations(a.iter(), 0, |_| 0).count());
512    assert_eq!(4, combinations(a.iter(), 1, |_| 0).count());
513    assert_eq!(6, combinations(a.iter(), 2, |_| 0).count());
514    assert_eq!(1, combinations(a.iter(), 4, |_| 0).count());
515    assert_eq!(0, combinations(a.iter(), 5, |_| 0).count());
516}
517
518#[test]
519fn test_powerset() {
520    let a = [1, 2, 3, 4];
521    assert_eq!(16, basic_powerset(a.iter()).count());
522
523    let a = [2, 3];
524    assert_eq!(vec![
525        vec![].into_boxed_slice(),
526        vec![2].into_boxed_slice(),
527        vec![3].into_boxed_slice(),
528        vec![2, 3].into_boxed_slice() 
529    ], basic_powerset(a.iter().copied()).collect::<Vec<_>>());
530
531    let a = [1, 2, 3];
532    assert_eq!(vec![
533        vec![].into_boxed_slice(),
534        vec![1].into_boxed_slice(),
535        vec![2].into_boxed_slice(),
536        vec![3].into_boxed_slice(),
537        vec![1, 2].into_boxed_slice(),
538        vec![1, 3].into_boxed_slice(),
539        vec![2, 3].into_boxed_slice(),
540        vec![1, 2, 3].into_boxed_slice()
541    ], basic_powerset(a.iter().copied()).collect::<Vec<_>>());
542}
543
544#[allow(deprecated)]
545#[test]
546fn test_multi_cartesian_product() {
547    let a = [0, 1];
548    let b = [0, 1];
549    let c = [-1, 1];
550    let all = [a, b, c];
551    let it = multi_cartesian_product(
552        all.iter().map(|l| l.iter().map(|x| *x)), 
553        |x| [x[0], x[1], x[2]],
554        |_, x| *x
555    );
556    let expected = vec![
557        [0, 0, -1],
558        [0, 0, 1],
559        [0, 1, -1],
560        [0, 1, 1],
561        [1, 0, -1],
562        [1, 0, 1],
563        [1, 1, -1],
564        [1, 1, 1]
565    ];
566    assert_eq!(expected, it.collect::<Vec<_>>());
567}
568
569#[allow(trivial_casts)]
570#[test]
571fn test_multiset_combinations() {
572    let a = [1, 2, 3, 1];
573    let mut iter = multiset_combinations(&a, 3, clone_slice);
574    assert_eq!(&[1, 2, 0, 0][..], &*iter.next().unwrap());
575    assert_eq!(&[1, 1, 1, 0][..], &*iter.next().unwrap());
576    assert_eq!(&[1, 1, 0, 1][..], &*iter.next().unwrap());
577    assert_eq!(&[1, 0, 2, 0][..], &*iter.next().unwrap());
578    assert_eq!(&[1, 0, 1, 1][..], &*iter.next().unwrap());
579
580    assert_eq!(&[0, 2, 1, 0][..], &*iter.next().unwrap());
581    assert_eq!(&[0, 2, 0, 1][..], &*iter.next().unwrap());
582    assert_eq!(&[0, 1, 2, 0][..], &*iter.next().unwrap());
583    assert_eq!(&[0, 1, 1, 1][..], &*iter.next().unwrap());
584
585    assert_eq!(&[0, 0, 3, 0][..], &*iter.next().unwrap());
586    assert_eq!(&[0, 0, 2, 1][..], &*iter.next().unwrap());
587    assert_eq!(None, iter.next());
588
589    assert_eq!(&([] as [[usize; 1]; 0]), &multiset_combinations(&[0], 1, clone_array::<usize, 1>).collect::<Vec<_>>()[..]);
590    assert_eq!(&([[0]] as [[usize; 1]; 1]), &multiset_combinations(&[0], 0, clone_array::<usize, 1>).collect::<Vec<_>>()[..]);
591
592    let a = [0, 2, 0, 1];
593    let mut iter = multiset_combinations(&a, 2, clone_slice);
594    assert_eq!(&[0, 2, 0, 0][..], &*iter.next().unwrap());
595    assert_eq!(&[0, 1, 0, 1][..], &*iter.next().unwrap());
596    assert_eq!(None, iter.next());
597
598    let a = [0, 3, 0, 0, 2, 0];
599    let mut iter = multiset_combinations(&a, 3, clone_slice);
600    assert_eq!(&[0, 3, 0, 0, 0, 0][..], &*iter.next().unwrap());
601    assert_eq!(&[0, 2, 0, 0, 1, 0][..], &*iter.next().unwrap());
602    assert_eq!(&[0, 1, 0, 0, 2, 0][..], &*iter.next().unwrap());
603    assert_eq!(None, iter.next());
604
605    let a = [0, 3, 0, 0, 2, 2, 0];
606    let mut iter = multiset_combinations(&a, 3, clone_slice);
607    assert_eq!(&[0, 3, 0, 0, 0, 0, 0][..], &*iter.next().unwrap());
608    assert_eq!(&[0, 2, 0, 0, 1, 0, 0][..], &*iter.next().unwrap());
609    assert_eq!(&[0, 2, 0, 0, 0, 1, 0][..], &*iter.next().unwrap());
610    assert_eq!(&[0, 1, 0, 0, 2, 0, 0][..], &*iter.next().unwrap());
611    assert_eq!(&[0, 1, 0, 0, 1, 1, 0][..], &*iter.next().unwrap());
612    assert_eq!(&[0, 1, 0, 0, 0, 2, 0][..], &*iter.next().unwrap());
613    assert_eq!(&[0, 0, 0, 0, 2, 1, 0][..], &*iter.next().unwrap());
614    assert_eq!(&[0, 0, 0, 0, 1, 2, 0][..], &*iter.next().unwrap());
615    assert_eq!(None, iter.next());
616
617    let a = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0];
618    let mut iter = multiset_combinations(&a, 10, clone_slice);
619    assert_eq!(&[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0][..], &*iter.next().unwrap());
620    assert_eq!(1000, iter.count());
621}
622
623#[test]
624fn test_multiset_combinations_k_unlimited() {
625    fn fac(n: usize) -> usize {
626        if n == 0 { 1 } else { n * fac(n - 1) }
627    }
628    let a = [10, 10, 10, 10, 10, 10];
629    assert_eq!(1, multiset_combinations(&a[..], 0, |_| ()).count());
630    assert_eq!(6, multiset_combinations(&a[..], 1, |_| ()).count());
631    assert_eq!(fac(6 + 8 - 1) / fac(6 - 1) / fac(8), multiset_combinations(&a[..], 8, |_| ()).count());
632}
feanor_math/iters.rs

feanor_math/
iters.rs
