Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Powerset count #735

Merged
merged 11 commits into from
Aug 26, 2023
11 changes: 8 additions & 3 deletions src/combinations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,12 @@ impl<I: Iterator> Combinations<I> {
self.pool.prefill(k);
}
}

pub(crate) fn n_and_count(self) -> (usize, usize) {
let Self { indices, pool, first } = self;
let n = pool.count();
(n, remaining_for(n, first, &indices).unwrap())
}
}

impl<I> Iterator for Combinations<I>
Expand Down Expand Up @@ -128,10 +134,9 @@ impl<I> Iterator for Combinations<I>
(low, upp)
}

#[inline]
fn count(self) -> usize {
let Self { indices, pool, first } = self;
let n = pool.count();
remaining_for(n, first, &indices).unwrap()
self.n_and_count().1
}
}

Expand Down
45 changes: 22 additions & 23 deletions src/powerset.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@ use std::iter::FusedIterator;
use std::usize;
use alloc::vec::Vec;

use super::combinations::{Combinations, combinations};
use super::size_hint;
use super::combinations::{Combinations, checked_binomial, combinations};

/// An iterator to iterate through the powerset of the elements from an iterator.
///
Expand All @@ -13,22 +12,20 @@ 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,
{
clone_fields!(combs, pos);
clone_fields!(combs);
}

impl<I> fmt::Debug for Powerset<I>
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.
Expand All @@ -38,7 +35,6 @@ pub fn powerset<I>(src: I) -> Powerset<I>
{
Powerset {
combs: combinations(src, 0),
pos: 0,
}
}

Expand All @@ -51,35 +47,32 @@ impl<I> Iterator for Powerset<I>

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
{
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>) {
let k = self.combs.k();
// Total bounds for source iterator.
let src_total = self.combs.src().size_hint();

// Total bounds for self ( length(powerset(set) == 2 ^ length(set) )
let self_total = size_hint::pow_scalar_base(2, src_total);
let (n_min, n_max) = self.combs.src().size_hint();
// Total bounds for the current combinations.
let (mut low, mut upp) = self.combs.size_hint();
low = remaining_for(low, n_min, k).unwrap_or(usize::MAX);
upp = upp.and_then(|upp| remaining_for(upp, n_max?, k));
Philippe-Cholet marked this conversation as resolved.
Show resolved Hide resolved
(low, upp)
}

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 count(self) -> usize {
let k = self.combs.k();
let (n, combs_count) = self.combs.n_and_count();
remaining_for(combs_count, n, k).unwrap()
}
}

Expand All @@ -88,3 +81,9 @@ impl<I> FusedIterator for Powerset<I>
I: Iterator,
I::Item: Clone,
{}

fn remaining_for(init_count: usize, n: usize, k: usize) -> Option<usize> {
Philippe-Cholet marked this conversation as resolved.
Show resolved Hide resolved
(k + 1..=n).fold(Some(init_count), |sum, i| {
sum.and_then(|s| s.checked_add(checked_binomial(n, i)?))
})
}
1 change: 1 addition & 0 deletions src/size_hint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ pub fn mul_scalar(sh: SizeHint, x: usize) -> SizeHint {

/// Raise `base` correctly by a `SizeHint` exponent.
#[inline]
#[allow(dead_code)]
Philippe-Cholet marked this conversation as resolved.
Show resolved Hide resolved
pub fn pow_scalar_base(base: usize, exp: SizeHint) -> SizeHint {
let exp_low = cmp::min(exp.0, u32::MAX as usize) as u32;
let low = base.saturating_pow(exp_low);
Expand Down
17 changes: 17 additions & 0 deletions tests/test_std.rs
Original file line number Diff line number Diff line change
Expand Up @@ -989,6 +989,23 @@ fn powerset() {
assert_eq!((0..4).powerset().count(), 1 << 4);
assert_eq!((0..8).powerset().count(), 1 << 8);
assert_eq!((0..16).powerset().count(), 1 << 16);

for n in 0..=10 {
let mut it = (0..n).powerset();
let len = 1 << n;
Philippe-Cholet marked this conversation as resolved.
Show resolved Hide resolved
assert_eq!(len, it.clone().count());
assert_eq!(len, it.size_hint().0);
assert_eq!(Some(len), it.size_hint().1);
for count in (0..len).rev() {
let elem = it.next();
assert!(elem.is_some());
assert_eq!(count, it.clone().count());
assert_eq!(count, it.size_hint().0);
assert_eq!(Some(count), it.size_hint().1);
}
let should_be_none = it.next();
assert!(should_be_none.is_none());
}
}

#[test]
Expand Down