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(iter:src, k: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(init:0usize, |sum: usize, i: usize| sum.checked_add(checked_binomial(n, k:i)?))
131}
132