Skip to content

Commit

Permalink
FEAT: Powerset iterator adaptor
Browse files Browse the repository at this point in the history
An iterator to iterate through the powerset of the elements from an iterator.
  • Loading branch information
willcrozi committed Sep 29, 2020
1 parent 2b005c2 commit 8cdf928
Show file tree
Hide file tree
Showing 6 changed files with 169 additions and 12 deletions.
46 changes: 38 additions & 8 deletions src/combinations.rs
Expand Up @@ -31,13 +31,8 @@ impl<I> fmt::Debug for Combinations<I>
pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
where I: Iterator
{
let mut pool: LazyBuffer<I> = LazyBuffer::new(iter);

for _ in 0..k {
if !pool.get_next() {
break;
}
}
let mut pool = LazyBuffer::new(iter);
pool.prefill(k);

Combinations {
indices: (0..k).collect(),
Expand All @@ -46,14 +41,49 @@ pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
}
}

impl<I: Iterator> Combinations<I> {
/// Returns the length of a combination produced by this iterator.
#[inline]
pub fn k(&self) -> usize { self.indices.len() }

/// Returns the (current) length of the pool from which combination elements are
/// selected. This value can change between invocations of [`next`].
///
/// [`next`]: #method.next
#[inline]
pub fn n(&self) -> usize { self.pool.len() }

/// Resets this `Combinations` back to an initial state for combinations of length
/// `k` over the same pool data source. If `k` is larger than the current length
/// of the data pool an attempt is made to prefill the pool so that it holds `k`
/// elements.
pub(crate) fn reset(&mut self, k: usize) {
self.first = true;

if k < self.indices.len() {
self.indices.truncate(k);
for i in 0..k {
self.indices[i] = i;
}

} else {
for i in 0..self.indices.len() {
self.indices[i] = i;
}
self.indices.extend(self.indices.len()..k);
self.pool.prefill(k);
}
}
}

impl<I> Iterator for Combinations<I>
where I: Iterator,
I::Item: Clone
{
type Item = Vec<I::Item>;
fn next(&mut self) -> Option<Self::Item> {
if self.first {
if self.pool.is_done() {
if self.k() > self.n() {
return None;
}
self.first = false;
Expand Down
15 changes: 11 additions & 4 deletions src/lazy_buffer.rs
Expand Up @@ -24,10 +24,6 @@ where
self.buffer.len()
}

pub fn is_done(&self) -> bool {
self.done
}

pub fn get_next(&mut self) -> bool {
if self.done {
return false;
Expand All @@ -44,6 +40,17 @@ where
}
}
}

pub fn prefill(&mut self, len: usize) {
let buffer_len = self.buffer.len();

if !self.done && len > buffer_len {
let delta = len - buffer_len;

self.buffer.extend(self.it.by_ref().take(delta));
self.done = self.buffer.len() < len;
}
}
}

impl<I, J> Index<J> for LazyBuffer<I>
Expand Down
40 changes: 40 additions & 0 deletions src/lib.rs
Expand Up @@ -134,6 +134,8 @@ pub mod structs {
pub use crate::permutations::Permutations;
pub use crate::process_results_impl::ProcessResults;
#[cfg(feature = "use_alloc")]
pub use crate::powerset::Powerset;
#[cfg(feature = "use_alloc")]
pub use crate::put_back_n_impl::PutBackN;
#[cfg(feature = "use_alloc")]
pub use crate::rciter_impl::RcIter;
Expand Down Expand Up @@ -207,6 +209,8 @@ mod peek_nth;
mod peeking_take_while;
#[cfg(feature = "use_alloc")]
mod permutations;
#[cfg(feature = "use_alloc")]
mod powerset;
mod process_results_impl;
#[cfg(feature = "use_alloc")]
mod put_back_n_impl;
Expand Down Expand Up @@ -1406,6 +1410,42 @@ pub trait Itertools : Iterator {
permutations::permutations(self, k)
}

/// Return an iterator that iterates through the powerset of the elements from an
/// iterator.
///
/// Iterator element type is `Vec<Self::Item>`. The iterator produces a new `Vec`
/// per iteration, and clones the iterator elements.
///
/// The powerset of a set contains all subsets including the empty set and the full
/// input set. A powerset has length _2^n_ where _n_ is the length of the input
/// set.
///
/// Each `Vec` produced by this iterator represents a subset of the elements
/// produced by the source iterator.
///
/// ```
/// use itertools::Itertools;
///
/// let sets = (1..4).powerset().collect::<Vec<_>>();
/// itertools::assert_equal(sets, vec![
/// vec![],
/// vec![1],
/// vec![2],
/// vec![3],
/// vec![1, 2],
/// vec![1, 3],
/// vec![2, 3],
/// vec![1, 2, 3],
/// ]);
/// ```
#[cfg(feature = "use_alloc")]
fn powerset(self) -> Powerset<Self>
where Self: Sized,
Self::Item: Clone,
{
powerset::powerset(self)
}

/// Return an iterator adaptor that pads the sequence to a minimum length of
/// `min` by filling missing elements using a closure `f`.
///
Expand Down
56 changes: 56 additions & 0 deletions src/powerset.rs
@@ -0,0 +1,56 @@
use std::fmt;
use alloc::vec::Vec;

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

/// An iterator to iterate through the powerset of the elements from an iterator.
///
/// See [`.powerset()`](../trait.Itertools.html#method.powerset) for more
/// information.
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct Powerset<I: Iterator> {
combs: Combinations<I>,
}

impl<I> Clone for Powerset<I>
where I: Clone + Iterator,
I::Item: Clone,
{
clone_fields!(combs);
}

impl<I> fmt::Debug for Powerset<I>
where I: Iterator + fmt::Debug,
I::Item: fmt::Debug,
{
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,
{
Powerset { combs: combinations(src, 0) }
}

impl<I> Iterator for Powerset<I>
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() {
Some(elt)
} else if self.combs.k() < self.combs.n()
|| self.combs.k() == 0
{
self.combs.reset(self.combs.k() + 1);
self.combs.next()
} else {
None
}
}
}
7 changes: 7 additions & 0 deletions tests/quick.rs
Expand Up @@ -907,6 +907,13 @@ quickcheck! {
}
}

quickcheck! {
fn size_powerset(it: Iter<u8, Exact>) -> bool {
// Powerset cardinality gets large very quickly, limit input to keep test fast.
correct_size_hint(it.take(12).powerset())
}
}

quickcheck! {
fn size_unique(it: Iter<i8>) -> bool {
correct_size_hint(it.unique())
Expand Down
17 changes: 17 additions & 0 deletions tests/test_std.rs
Expand Up @@ -764,6 +764,23 @@ fn combinations_with_replacement() {
);
}

#[test]
fn powerset() {
it::assert_equal((0..0).powerset(), vec![vec![]]);
it::assert_equal((0..1).powerset(), vec![vec![], vec![0]]);
it::assert_equal((0..2).powerset(), vec![vec![], vec![0], vec![1], vec![0, 1]]);
it::assert_equal((0..3).powerset(), vec![
vec![],
vec![0], vec![1], vec![2],
vec![0, 1], vec![0, 2], vec![1, 2],
vec![0, 1, 2]
]);

assert_eq!((0..4).powerset().count(), 1 << 4);
assert_eq!((0..8).powerset().count(), 1 << 8);
assert_eq!((0..16).powerset().count(), 1 << 16);
}

#[test]
fn diff_mismatch() {
let a = vec![1, 2, 3, 4];
Expand Down

0 comments on commit 8cdf928

Please sign in to comment.