Skip to content

Add array_combinations_with_replacement #1033

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

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
31 changes: 31 additions & 0 deletions benches/combinations_with_replacement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,42 @@ fn comb_replacement_n10_k10(c: &mut Criterion) {
})
});
}
fn array_comb_replacement_n10_k5(c: &mut Criterion) {
c.bench_function("array comb replacement n10k5", move |b| {
b.iter(|| {
for i in (0..10).array_combinations_with_replacement::<5>() {
black_box(i);
}
})
});
}

fn array_comb_replacement_n5_k10(c: &mut Criterion) {
c.bench_function("array comb replacement n5 k10", move |b| {
b.iter(|| {
for i in (0..5).array_combinations_with_replacement::<10>() {
black_box(i);
}
})
});
}

fn array_comb_replacement_n10_k10(c: &mut Criterion) {
c.bench_function("array comb replacement n10 k10", move |b| {
b.iter(|| {
for i in (0..10).array_combinations_with_replacement::<10>() {
black_box(i);
}
})
});
}
criterion_group!(
benches,
comb_replacement_n10_k5,
comb_replacement_n5_k10,
comb_replacement_n10_k10,
array_comb_replacement_n10_k5,
array_comb_replacement_n5_k10,
array_comb_replacement_n10_k10,
);
criterion_main!(benches);
24 changes: 24 additions & 0 deletions benches/specializations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -441,6 +441,30 @@ bench_specializations! {
}
v.iter().combinations_with_replacement(4)
}
array_combinations_with_replacement1 {
{
let v = black_box(vec![0; 4096]);
}
v.iter().array_combinations_with_replacement::<1>()
}
array_combinations_with_replacement2 {
{
let v = black_box(vec![0; 90]);
}
v.iter().array_combinations_with_replacement::<2>()
}
array_combinations_with_replacement3 {
{
let v = black_box(vec![0; 28]);
}
v.iter().array_combinations_with_replacement::<3>()
}
array_combinations_with_replacement4 {
{
let v = black_box(vec![0; 16]);
}
v.iter().array_combinations_with_replacement::<4>()
}
permutations1 {
{
let v = black_box(vec![0; 1024]);
Expand Down
10 changes: 10 additions & 0 deletions src/combinations.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use alloc::boxed::Box;
use core::array;
use core::borrow::BorrowMut;
use std::fmt;
Expand Down Expand Up @@ -52,7 +53,16 @@ pub trait PoolIndex<T>: BorrowMut<[usize]> {
self.borrow().len()
}
}
impl<T> PoolIndex<T> for Box<[usize]> {
type Item = Vec<T>;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can think about returning Box<[T]> later, but since we'd probably construct it via a Vec, Vec is fine.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just what did not change that, as that was the way it was before.


fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
where
T: Clone,
{
pool.get_at(self)
}
}
impl<T> PoolIndex<T> for Vec<usize> {
type Item = Vec<T>;

Expand Down
71 changes: 46 additions & 25 deletions src/combinations_with_replacement.rs
Original file line number Diff line number Diff line change
@@ -1,52 +1,62 @@
use alloc::boxed::Box;
use alloc::vec::Vec;
use std::fmt;
use std::iter::FusedIterator;

use super::lazy_buffer::LazyBuffer;
use crate::adaptors::checked_binomial;

use crate::combinations::PoolIndex;
/// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement.
///
/// See [`.combinations_with_replacement()`](crate::Itertools::combinations_with_replacement)
/// for more information.
#[derive(Clone)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct CombinationsWithReplacement<I>
pub struct CombinationsWithReplacementGeneric<I, Idx>
where
I: Iterator,
I::Item: Clone,
{
indices: Box<[usize]>,
indices: Idx,
pool: LazyBuffer<I>,
first: bool,
}

impl<I> fmt::Debug for CombinationsWithReplacement<I>
/// Iterator for `Box<[I]>` valued combinations_with_replacement returned by [`.combinations_with_replacement()`](crate::Itertools::combinations_with_replacement)
pub type CombinationsWithReplacement<I> = CombinationsWithReplacementGeneric<I, Box<[usize]>>;
/// Iterator for const generic combinations_with_replacement returned by [`.array_combinations_with_replacement()`](crate::Itertools::array_combinations_with_replacement)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-markdown expert here: Is the (crate::Itertools::array_combinations_with_replacement) required?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know I just copied and modified the doc comment from two lines above.

pub type ArrayCombinationsWithReplacement<I, const K: usize> =
CombinationsWithReplacementGeneric<I, [usize; K]>;

impl<I, Idx> fmt::Debug for CombinationsWithReplacementGeneric<I, Idx>
where
I: Iterator + fmt::Debug,
I::Item: fmt::Debug + Clone,
Idx: fmt::Debug,
{
debug_fmt_fields!(CombinationsWithReplacement, indices, pool, first);
debug_fmt_fields!(CombinationsWithReplacementGeneric, indices, pool, first);
}

/// Create a new `ArrayCombinationsWithReplacement`` from a clonable iterator.
pub fn array_combinations_with_replacement<I: Iterator, const K: usize>(
iter: I,
) -> ArrayCombinationsWithReplacement<I, K>
where
I::Item: Clone,
{
ArrayCombinationsWithReplacement::new(iter, [0; K])
}
/// Create a new `CombinationsWithReplacement` from a clonable iterator.
pub fn combinations_with_replacement<I>(iter: I, k: usize) -> CombinationsWithReplacement<I>
where
I: Iterator,
I::Item: Clone,
{
let indices = alloc::vec![0; k].into_boxed_slice();
let pool: LazyBuffer<I> = LazyBuffer::new(iter);

CombinationsWithReplacement {
indices,
pool,
first: true,
}
CombinationsWithReplacementGeneric::new(iter, indices)
}

impl<I> CombinationsWithReplacement<I>
impl<I: Iterator, Idx: PoolIndex<I::Item>> CombinationsWithReplacementGeneric<I, Idx>
where
I: Iterator,
I::Item: Clone,
Expand All @@ -62,7 +72,8 @@ where

// Work out where we need to update our indices
let mut increment = None;
for (i, indices_int) in self.indices.iter().enumerate().rev() {
let indices: &mut [usize] = self.indices.borrow_mut();
for (i, indices_int) in indices.iter().enumerate().rev() {
if *indices_int < self.pool.len() - 1 {
increment = Some((i, indices_int + 1));
break;
Expand All @@ -73,39 +84,48 @@ where
Some((increment_from, increment_value)) => {
// We need to update the rightmost non-max value
// and all those to the right
self.indices[increment_from..].fill(increment_value);
indices[increment_from..].fill(increment_value);
false
}
// Otherwise, we're done
None => true,
}
}
/// Constructor with arguments the inner iterator and the initial state for the indices.
fn new(iter: I, indices: Idx) -> Self {
Self {
indices,
pool: LazyBuffer::new(iter),
first: true,
}
}
}

impl<I> Iterator for CombinationsWithReplacement<I>
impl<I, Idx> Iterator for CombinationsWithReplacementGeneric<I, Idx>
where
I: Iterator,
I::Item: Clone,
Idx: PoolIndex<I::Item>,
{
type Item = Vec<I::Item>;
type Item = Idx::Item;

fn next(&mut self) -> Option<Self::Item> {
if self.first {
// In empty edge cases, stop iterating immediately
if !(self.indices.is_empty() || self.pool.get_next()) {
if !(self.indices.borrow().is_empty() || self.pool.get_next()) {
return None;
}
self.first = false;
} else if self.increment_indices() {
return None;
}
Some(self.pool.get_at(&self.indices))
Some(self.indices.extract_item(&self.pool))
}

fn nth(&mut self, n: usize) -> Option<Self::Item> {
if self.first {
// In empty edge cases, stop iterating immediately
if !(self.indices.is_empty() || self.pool.get_next()) {
if !(self.indices.borrow().is_empty() || self.pool.get_next()) {
return None;
}
self.first = false;
Expand All @@ -117,13 +137,13 @@ where
return None;
}
}
Some(self.pool.get_at(&self.indices))
Some(self.indices.extract_item(&self.pool))
}

fn size_hint(&self) -> (usize, Option<usize>) {
let (mut low, mut upp) = self.pool.size_hint();
low = remaining_for(low, self.first, &self.indices).unwrap_or(usize::MAX);
upp = upp.and_then(|upp| remaining_for(upp, self.first, &self.indices));
low = remaining_for(low, self.first, self.indices.borrow()).unwrap_or(usize::MAX);
upp = upp.and_then(|upp| remaining_for(upp, self.first, self.indices.borrow()));
(low, upp)
}

Expand All @@ -134,14 +154,15 @@ where
first,
} = self;
let n = pool.count();
remaining_for(n, first, &indices).unwrap()
remaining_for(n, first, indices.borrow()).unwrap()
}
}

impl<I> FusedIterator for CombinationsWithReplacement<I>
impl<I, Idx> FusedIterator for CombinationsWithReplacementGeneric<I, Idx>
where
I: Iterator,
I::Item: Clone,
Idx: PoolIndex<I::Item>,
{
}

Expand Down
32 changes: 31 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,8 @@ pub mod traits {
pub use crate::tuple_impl::HomogeneousTuple;
}

#[cfg(feature = "use_alloc")]
use crate::combinations_with_replacement::ArrayCombinationsWithReplacement;
pub use crate::concat_impl::concat;
pub use crate::cons_tuples_impl::cons_tuples;
pub use crate::diff::diff_with;
Expand Down Expand Up @@ -1804,7 +1806,35 @@ pub trait Itertools: Iterator {
{
combinations_with_replacement::combinations_with_replacement(self, k)
}

/// Return an iterator that iterates over the `k`-length combinations of
/// the elements from an iterator, with replacement.
///
/// Iterator element type is [Self::Item; K]. The iterator produces a new
/// array per iteration, and clones the iterator elements.
///
/// ```
/// use itertools::Itertools;
///
/// let it = (1..4).array_combinations_with_replacement::<2>();
/// itertools::assert_equal(it, vec![
/// [1, 1],
/// [1, 2],
/// [1, 3],
/// [2, 2],
/// [2, 3],
/// [3, 3],
/// ]);
/// ```
#[cfg(feature = "use_alloc")]
fn array_combinations_with_replacement<const K: usize>(
self,
) -> ArrayCombinationsWithReplacement<Self, K>
where
Self: Sized,
Self::Item: Clone,
{
combinations_with_replacement::array_combinations_with_replacement(self)
}
/// Return an iterator adaptor that iterates over all k-permutations of the
/// elements from an iterator.
///
Expand Down
4 changes: 4 additions & 0 deletions tests/adaptors_no_collect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,3 +49,7 @@ fn combinations_no_collect() {
fn combinations_with_replacement_no_collect() {
no_collect_test(|iter| iter.combinations_with_replacement(5))
}
#[test]
fn array_combinations_with_replacement_no_collect() {
no_collect_test(|iter| iter.array_combinations_with_replacement::<5>())
}
5 changes: 5 additions & 0 deletions tests/laziness.rs
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,11 @@ must_use_tests! {
let _ = Panicking.combinations_with_replacement(1);
let _ = Panicking.combinations_with_replacement(2);
}
array_combinations_with_replacement {
let _ = Panicking.array_combinations_with_replacement::<0>();
let _ = Panicking.array_combinations_with_replacement::<1>();
let _ = Panicking.array_combinations_with_replacement::<2>();
}
permutations {
let _ = Panicking.permutations(0);
let _ = Panicking.permutations(1);
Expand Down
5 changes: 5 additions & 0 deletions tests/quick.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1838,6 +1838,11 @@ quickcheck! {
is_fused(a.combinations_with_replacement(3))
}

fn fused_array_combination_with_replacement(a: Iter<i16>) -> bool
{
is_fused(a.clone().array_combinations_with_replacement::<1>()) &&
is_fused(a.array_combinations_with_replacement::<3>())
}
fn fused_tuple_combination(a: Iter<i16>) -> bool
{
is_fused(a.clone().fuse().tuple_combinations::<(_,)>()) &&
Expand Down
10 changes: 10 additions & 0 deletions tests/specializations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,16 @@ quickcheck! {
TestResult::passed()
}

fn array_combinations_with_replacement(a: Vec<u8>) -> TestResult {
if a.len() > 10 {
return TestResult::discard();
}
test_specializations(&a.iter().array_combinations_with_replacement::<1>());
test_specializations(&a.iter().array_combinations_with_replacement::<2>());
test_specializations(&a.iter().array_combinations_with_replacement::<3>());

TestResult::passed()
}
fn permutations(a: Vec<u8>, n: u8) -> TestResult {
if n > 3 || a.len() > 8 {
return TestResult::discard();
Expand Down
24 changes: 24 additions & 0 deletions tests/test_std.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1256,6 +1256,30 @@ fn combinations_with_replacement_range_count() {
}
}

#[test]
#[cfg(not(miri))]
fn array_combinations_with_replacement() {
// Pool smaller than n
it::assert_equal(
(0..1).array_combinations_with_replacement::<2>(),
vec![[0, 0]],
);
// Pool larger than n
it::assert_equal(
(0..3).array_combinations_with_replacement::<2>(),
vec![[0, 0], [0, 1], [0, 2], [1, 1], [1, 2], [2, 2]],
);
// Zero size
it::assert_equal((0..3).array_combinations_with_replacement::<0>(), vec![[]]);
// Zero size on empty pool
it::assert_equal((0..0).array_combinations_with_replacement::<0>(), vec![[]]);
// Empty pool
it::assert_equal(
(0..0).array_combinations_with_replacement::<2>(),
vec![] as Vec<[_; 2]>,
);
}

#[test]
fn powerset() {
it::assert_equal((0..0).powerset(), vec![vec![]]);
Expand Down