diff options
Diffstat (limited to 'vendor/itertools/src/powerset.rs')
-rw-r--r-- | vendor/itertools/src/powerset.rs | 95 |
1 files changed, 56 insertions, 39 deletions
diff --git a/vendor/itertools/src/powerset.rs b/vendor/itertools/src/powerset.rs index 4d7685b12..9a7131a3e 100644 --- a/vendor/itertools/src/powerset.rs +++ b/vendor/itertools/src/powerset.rs @@ -1,10 +1,11 @@ +use alloc::vec::Vec; use std::fmt; use std::iter::FusedIterator; use std::usize; -use alloc::vec::Vec; -use super::combinations::{Combinations, combinations}; -use super::size_hint; +use super::combinations::{combinations, Combinations}; +use crate::adaptors::checked_binomial; +use crate::size_hint::{self, SizeHint}; /// An iterator to iterate through the powerset of the elements from an iterator. /// @@ -13,78 +14,94 @@ use super::size_hint; #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] pub struct Powerset<I: Iterator> { combs: Combinations<I>, - // Iterator `position` (equal to count of yielded elements). - pos: usize, } impl<I> Clone for Powerset<I> - where I: Clone + Iterator, - I::Item: Clone, +where + I: Clone + Iterator, + I::Item: Clone, { - clone_fields!(combs, pos); + clone_fields!(combs); } impl<I> fmt::Debug for Powerset<I> - where I: Iterator + fmt::Debug, - I::Item: fmt::Debug, +where + I: Iterator + fmt::Debug, + I::Item: fmt::Debug, { - debug_fmt_fields!(Powerset, combs, pos); + debug_fmt_fields!(Powerset, combs); } /// Create a new `Powerset` from a clonable iterator. pub fn powerset<I>(src: I) -> Powerset<I> - where I: Iterator, - I::Item: Clone, +where + I: Iterator, + I::Item: Clone, { Powerset { combs: combinations(src, 0), - pos: 0, } } impl<I> Iterator for Powerset<I> - where - I: Iterator, - I::Item: Clone, +where + I: Iterator, + I::Item: Clone, { type Item = Vec<I::Item>; fn next(&mut self) -> Option<Self::Item> { if let Some(elt) = self.combs.next() { - self.pos = self.pos.saturating_add(1); Some(elt) - } else if self.combs.k() < self.combs.n() - || self.combs.k() == 0 - { + } else if self.combs.k() < self.combs.n() || self.combs.k() == 0 { self.combs.reset(self.combs.k() + 1); - self.combs.next().map(|elt| { - self.pos = self.pos.saturating_add(1); - elt - }) + self.combs.next() } else { None } } - fn size_hint(&self) -> (usize, Option<usize>) { + fn size_hint(&self) -> SizeHint { + let k = self.combs.k(); // Total bounds for source iterator. - let src_total = size_hint::add_scalar(self.combs.src().size_hint(), self.combs.n()); + let (n_min, n_max) = self.combs.src().size_hint(); + let low = remaining_for(n_min, k).unwrap_or(usize::MAX); + let upp = n_max.and_then(|n| remaining_for(n, k)); + size_hint::add(self.combs.size_hint(), (low, upp)) + } - // Total bounds for self ( length(powerset(set) == 2 ^ length(set) ) - let self_total = size_hint::pow_scalar_base(2, src_total); + fn count(self) -> usize { + let k = self.combs.k(); + let (n, combs_count) = self.combs.n_and_count(); + combs_count + remaining_for(n, k).unwrap() + } - if self.pos < usize::MAX { - // Subtract count of elements already yielded from total. - size_hint::sub_scalar(self_total, self.pos) - } else { - // Fallback: self.pos is saturated and no longer reliable. - (0, self_total.1) + fn fold<B, F>(self, mut init: B, mut f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + let mut it = self.combs; + if it.k() == 0 { + init = it.by_ref().fold(init, &mut f); + it.reset(1); } + init = it.by_ref().fold(init, &mut f); + // n is now known for sure because k >= 1 and all k-combinations have been generated. + for k in it.k() + 1..=it.n() { + it.reset(k); + init = it.by_ref().fold(init, &mut f); + } + init } } impl<I> FusedIterator for Powerset<I> - where - I: Iterator, - I::Item: Clone, -{} +where + I: Iterator, + I::Item: Clone, +{ +} + +fn remaining_for(n: usize, k: usize) -> Option<usize> { + (k + 1..=n).try_fold(0usize, |sum, i| sum.checked_add(checked_binomial(n, i)?)) +} |