1 use alloc::vec::Vec;
2 use std::fmt;
3 use std::iter::FusedIterator;
4
5 use super::combinations::{combinations, Combinations};
6 use crate::adaptors::checked_binomial;
7 use 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"]
14 pub struct Powerset<I: Iterator> {
15 combs: Combinations<I>,
16 }
17
18 impl<I> Clone for Powerset<I>
19 where
20 I: Clone + Iterator,
21 I::Item: Clone,
22 {
23 clone_fields!(combs);
24 }
25
26 impl<I> fmt::Debug for Powerset<I>
27 where
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.
powerset<I>(src: I) -> Powerset<I> where I: Iterator, I::Item: Clone,35 pub fn powerset<I>(src: I) -> Powerset<I>
36 where
37 I: Iterator,
38 I::Item: Clone,
39 {
40 Powerset {
41 combs: combinations(src, 0),
42 }
43 }
44
45 impl<I: Iterator> Powerset<I> {
46 /// Returns true if `k` has been incremented, false otherwise.
increment_k(&mut self) -> bool47 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
57 impl<I> Iterator for Powerset<I>
58 where
59 I: Iterator,
60 I::Item: Clone,
61 {
62 type Item = Vec<I::Item>;
63
next(&mut self) -> Option<Self::Item>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
nth(&mut self, mut n: usize) -> Option<Self::Item>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
size_hint(&self) -> SizeHint88 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
count(self) -> usize97 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
fold<B, F>(self, mut init: B, mut f: F) -> B where F: FnMut(B, Self::Item) -> B,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
122 impl<I> FusedIterator for Powerset<I>
123 where
124 I: Iterator,
125 I::Item: Clone,
126 {
127 }
128
remaining_for(n: usize, k: usize) -> Option<usize>129 fn 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 }
132