-
Notifications
You must be signed in to change notification settings - Fork 322
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
base: master
Are you sure you want to change the base?
Changes from all commits
fbfa003
76dcd1e
a1dd959
bbbce31
9ce5068
189c15e
6607b1c
74b3157
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
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) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Non-markdown expert here: Is the There was a problem hiding this comment. Choose a reason for hiding this commentThe 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, | ||
|
@@ -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; | ||
|
@@ -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; | ||
|
@@ -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) | ||
} | ||
|
||
|
@@ -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>, | ||
{ | ||
} | ||
|
||
|
There was a problem hiding this comment.
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 aVec
,Vec
is fine.There was a problem hiding this comment.
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.