itertools/
powerset.rs

1use alloc::vec::Vec;
2use std::fmt;
3use std::iter::FusedIterator;
4use std::usize;
5
6use super::combinations::{combinations, Combinations};
7use crate::adaptors::checked_binomial;
8use crate::size_hint::{self, SizeHint};
9
10/// An iterator to iterate through the powerset of the elements from an iterator.
11///
12/// See [`.powerset()`](crate::Itertools::powerset) for more
13/// information.
14#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
15pub struct Powerset<I: Iterator> {
16    combs: Combinations<I>,
17}
18
19impl<I> Clone for Powerset<I>
20where
21    I: Clone + Iterator,
22    I::Item: Clone,
23{
24    clone_fields!(combs);
25}
26
27impl<I> fmt::Debug for Powerset<I>
28where
29    I: Iterator + fmt::Debug,
30    I::Item: fmt::Debug,
31{
32    debug_fmt_fields!(Powerset, combs);
33}
34
35/// Create a new `Powerset` from a clonable iterator.
36pub fn powerset<I>(src: I) -> Powerset<I>
37where
38    I: Iterator,
39    I::Item: Clone,
40{
41    Powerset {
42        combs: combinations(src, 0),
43    }
44}
45
46impl<I> Iterator for Powerset<I>
47where
48    I: Iterator,
49    I::Item: Clone,
50{
51    type Item = Vec<I::Item>;
52
53    fn next(&mut self) -> Option<Self::Item> {
54        if let Some(elt) = self.combs.next() {
55            Some(elt)
56        } else if self.combs.k() < self.combs.n() || self.combs.k() == 0 {
57            self.combs.reset(self.combs.k() + 1);
58            self.combs.next()
59        } else {
60            None
61        }
62    }
63
64    fn size_hint(&self) -> SizeHint {
65        let k = self.combs.k();
66        // Total bounds for source iterator.
67        let (n_min, n_max) = self.combs.src().size_hint();
68        let low = remaining_for(n_min, k).unwrap_or(usize::MAX);
69        let upp = n_max.and_then(|n| remaining_for(n, k));
70        size_hint::add(self.combs.size_hint(), (low, upp))
71    }
72
73    fn count(self) -> usize {
74        let k = self.combs.k();
75        let (n, combs_count) = self.combs.n_and_count();
76        combs_count + remaining_for(n, k).unwrap()
77    }
78
79    fn fold<B, F>(self, mut init: B, mut f: F) -> B
80    where
81        F: FnMut(B, Self::Item) -> B,
82    {
83        let mut it = self.combs;
84        if it.k() == 0 {
85            init = it.by_ref().fold(init, &mut f);
86            it.reset(1);
87        }
88        init = it.by_ref().fold(init, &mut f);
89        // n is now known for sure because k >= 1 and all k-combinations have been generated.
90        for k in it.k() + 1..=it.n() {
91            it.reset(k);
92            init = it.by_ref().fold(init, &mut f);
93        }
94        init
95    }
96}
97
98impl<I> FusedIterator for Powerset<I>
99where
100    I: Iterator,
101    I::Item: Clone,
102{
103}
104
105fn remaining_for(n: usize, k: usize) -> Option<usize> {
106    (k + 1..=n).try_fold(0usize, |sum, i| sum.checked_add(checked_binomial(n, i)?))
107}