| 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 | |