feanor_math/
iters.rs

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

feanor_math/
iters.rs
