From b3a92a91aad51c075dfea116ac6779e8bfe7b15d Mon Sep 17 00:00:00 2001 From: Ronno Das Date: Tue, 10 Sep 2024 15:32:14 +0200 Subject: [PATCH] array_combinations using array::from_fn --- src/array_combinations.rs | 166 ++++++++++++++++++++++++++++++++++++++ src/combinations.rs | 14 +++- src/lazy_buffer.rs | 4 + src/lib.rs | 47 +++++++++++ 4 files changed, 228 insertions(+), 3 deletions(-) create mode 100644 src/array_combinations.rs diff --git a/src/array_combinations.rs b/src/array_combinations.rs new file mode 100644 index 000000000..9de275234 --- /dev/null +++ b/src/array_combinations.rs @@ -0,0 +1,166 @@ +use core::iter::FusedIterator; +use core::{array, fmt}; + +use crate::combinations::{n_and_count, remaining_for}; +use crate::lazy_buffer::LazyBuffer; + +/// An iterator to iterate through all combinations in an iterator of `Clone`-able items that +/// produces arrays of a specific size. +/// +/// See [`.array_combinations()`](crate::Itertools::array_combinations) for more +/// information. +#[derive(Clone)] +#[must_use = "this iterator adaptor is not lazy but does nearly nothing unless consumed"] +pub struct ArrayCombinations +where + I::Item: Clone, +{ + indices: [usize; K], + pool: LazyBuffer, + first: bool, +} + +/// Create a new `ArrayCombinations` from a clonable iterator. +pub fn array_combinations(iter: I) -> ArrayCombinations +where + I::Item: Clone, +{ + let indices = array::from_fn(|i| i); + let pool = LazyBuffer::new(iter); + + ArrayCombinations { + indices, + pool, + first: true, + } +} + +impl ArrayCombinations +where + I::Item: Clone, +{ + /// Returns the (current) length of the pool from which combination elements are + /// selected. This value can change between invocations of [`next`](Combinations::next). + #[inline] + pub fn n(&self) -> usize { + self.pool.len() + } + + /// Initialises the iterator by filling a buffer with elements from the + /// iterator. Returns true if there are no combinations, false otherwise. + fn init(&mut self) -> bool { + self.pool.prefill(K); + let done = K > self.n(); + if !done { + self.first = false; + } + + done + } + + /// Increments indices representing the combination to advance to the next + /// (in lexicographic order by increasing sequence) combination. For example + /// if we have n=4 & k=2 then `[0, 1] -> [0, 2] -> [0, 3] -> [1, 2] -> ...` + /// + /// Returns true if we've run out of combinations, false otherwise. + fn increment_indices(&mut self) -> bool { + if K == 0 { + return true; // Done + } + + // Scan from the end, looking for an index to increment + let mut i: usize = K - 1; + + // Check if we need to consume more from the iterator + if self.indices[i] == self.pool.len() - 1 { + _ = self.pool.get_next(); // may change pool size + } + + while self.indices[i] == i + self.pool.len() - K { + if i > 0 { + i -= 1; + } else { + // Reached the last combination + return true; + } + } + + // Increment index, and reset the ones to its right + self.indices[i] += 1; + for j in i + 1..K { + self.indices[j] = self.indices[j - 1] + 1; + } + + // If we've made it this far, we haven't run out of combos + false + } + + /// Returns the n-th item or the number of successful steps. + pub(crate) fn try_nth(&mut self, n: usize) -> Result<::Item, usize> + where + I::Item: Clone, + { + let done = if self.first { + self.init() + } else { + self.increment_indices() + }; + if done { + return Err(0); + } + for i in 0..n { + if self.increment_indices() { + return Err(i + 1); + } + } + Ok(self.pool.get_array(self.indices)) + } +} + +impl Iterator for ArrayCombinations +where + I::Item: Clone, +{ + type Item = [I::Item; K]; + + fn next(&mut self) -> Option { + let done = if self.first { + self.init() + } else { + self.increment_indices() + }; + + (!done).then(|| self.pool.get_array(self.indices)) + } + + fn nth(&mut self, n: usize) -> Option { + self.try_nth(n).ok() + } + + fn size_hint(&self) -> (usize, Option) { + 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, upp) + } + + #[inline] + fn count(self) -> usize { + n_and_count(self.pool, self.first, &self.indices).1 + } +} + +impl fmt::Debug for ArrayCombinations +where + I: Iterator + fmt::Debug, + I::Item: Clone + fmt::Debug, +{ + debug_fmt_fields!(ArrayCombinations, indices, pool, first); +} + +impl FusedIterator for ArrayCombinations +where + I: FusedIterator, + I::Item: Clone, +{ +} diff --git a/src/combinations.rs b/src/combinations.rs index 6bb2f3ec6..423ffa85b 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -91,8 +91,7 @@ impl Combinations { pool, first, } = self; - let n = pool.count(); - (n, remaining_for(n, first, &indices).unwrap()) + n_and_count(pool, first, &indices) } /// Initialises the iterator by filling a buffer with elements from the @@ -210,8 +209,17 @@ where { } +pub(crate) fn n_and_count( + pool: LazyBuffer, + first: bool, + indices: &[usize], +) -> (usize, usize) { + let n = pool.count(); + (n, remaining_for(n, first, indices).unwrap()) +} + /// For a given size `n`, return the count of remaining combinations or None if it would overflow. -fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option { +pub(crate) fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option { let k = indices.len(); if n < k { Some(0) diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index fefcff8f5..fafa5f726 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -59,6 +59,10 @@ where pub fn get_at(&self, indices: &[usize]) -> Vec { indices.iter().map(|i| self.buffer[*i].clone()).collect() } + + pub fn get_array(&self, indices: [usize; K]) -> [I::Item; K] { + indices.map(|i| self.buffer[i].clone()) + } } impl Index for LazyBuffer diff --git a/src/lib.rs b/src/lib.rs index 604529a59..0a324ce04 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -96,6 +96,7 @@ pub mod structs { FilterOk, Interleave, InterleaveShortest, MapInto, MapOk, Positions, Product, PutBack, TakeWhileRef, TupleCombinations, Update, WhileSome, }; + pub use crate::array_combinations::ArrayCombinations; #[cfg(feature = "use_alloc")] pub use crate::combinations::Combinations; #[cfg(feature = "use_alloc")] @@ -177,6 +178,7 @@ pub use crate::either_or_both::EitherOrBoth; pub mod free; #[doc(inline)] pub use crate::free::*; +mod array_combinations; #[cfg(feature = "use_alloc")] mod combinations; #[cfg(feature = "use_alloc")] @@ -1671,6 +1673,51 @@ pub trait Itertools: Iterator { adaptors::tuple_combinations(self) } + /// Return an iterator adaptor that iterates over the combinations of the + /// elements from an iterator. + /// + /// Iterator element can be any array of type `Self::Item`. + /// + /// # Guarantees + /// + /// If the adapted iterator is deterministic, + /// this iterator adapter yields items in a reliable order. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let mut v = Vec::new(); + /// for [a, b] in (1..5).array_combinations() { + /// v.push([a, b]); + /// } + /// assert_eq!(v, vec![[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]); + /// + /// let mut it = (1..5).array_combinations(); + /// assert_eq!(Some([1, 2, 3]), it.next()); + /// assert_eq!(Some([1, 2, 4]), it.next()); + /// assert_eq!(Some([1, 3, 4]), it.next()); + /// assert_eq!(Some([2, 3, 4]), it.next()); + /// assert_eq!(None, it.next()); + /// + /// // this requires a type hint + /// let it = (1..5).array_combinations::<3>(); + /// itertools::assert_equal(it, vec![[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]); + /// + /// // you can also specify the complete type + /// use itertools::ArrayCombinations; + /// use std::ops::Range; + /// + /// let it: ArrayCombinations, 3> = (1..5).array_combinations(); + /// itertools::assert_equal(it, vec![[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]); + /// ``` + fn array_combinations(self) -> ArrayCombinations + where + Self: Sized + Clone, + Self::Item: Clone, + { + array_combinations::array_combinations(self) + } + /// Return an iterator adaptor that iterates over the `k`-length combinations of /// the elements from an iterator. ///