1 | use alloc::vec::Vec; |
2 | use std::fmt; |
3 | use std::iter::FusedIterator; |
4 | use std::usize; |
5 | |
6 | use super::combinations::{combinations, Combinations}; |
7 | use crate::adaptors::checked_binomial; |
8 | use 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" ] |
15 | pub struct Powerset<I: Iterator> { |
16 | combs: Combinations<I>, |
17 | } |
18 | |
19 | impl<I> Clone for Powerset<I> |
20 | where |
21 | I: Clone + Iterator, |
22 | I::Item: Clone, |
23 | { |
24 | clone_fields!(combs); |
25 | } |
26 | |
27 | impl<I> fmt::Debug for Powerset<I> |
28 | where |
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. |
36 | pub fn powerset<I>(src: I) -> Powerset<I> |
37 | where |
38 | I: Iterator, |
39 | I::Item: Clone, |
40 | { |
41 | Powerset { |
42 | combs: combinations(iter:src, k:0), |
43 | } |
44 | } |
45 | |
46 | impl<I> Iterator for Powerset<I> |
47 | where |
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 | |
98 | impl<I> FusedIterator for Powerset<I> |
99 | where |
100 | I: Iterator, |
101 | I::Item: Clone, |
102 | { |
103 | } |
104 | |
105 | fn remaining_for(n: usize, k: usize) -> Option<usize> { |
106 | (k + 1..=n).try_fold(init:0usize, |sum: usize, i: usize| sum.checked_add(checked_binomial(n, k:i)?)) |
107 | } |
108 | |