itertools/
powerset.rs

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

itertools/
powerset.rs
