diff --git a/Cargo.toml b/Cargo.toml index 46b209417..4b2b2eb10 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ test = false [dependencies] either = { version = "1.0", default-features = false } +array-init = "2.0" [dev-dependencies] rand = "0.7" diff --git a/benches/tuple_combinations.rs b/benches/tuple_combinations.rs index 4e26b282e..a9255bf83 100644 --- a/benches/tuple_combinations.rs +++ b/benches/tuple_combinations.rs @@ -99,6 +99,51 @@ fn tuple_comb_c4(c: &mut Criterion) { }); } + +fn array_comb_c1(c: &mut Criterion) { + c.bench_function("array comb c1", move |b| { + b.iter(|| { + for [i] in (0..N1).array_combinations() { + black_box(i); + } + }) + }); +} + + +fn array_comb_c2(c: &mut Criterion) { + c.bench_function("array comb c2", move |b| { + b.iter(|| { + for [i, j] in (0..N2).array_combinations() { + black_box(i + j); + } + }) + }); +} + + +fn array_comb_c3(c: &mut Criterion) { + c.bench_function("array comb c3", move |b| { + b.iter(|| { + for [i, j, k] in (0..N3).array_combinations() { + black_box(i + j + k); + } + }) + }); +} + + +fn array_comb_c4(c: &mut Criterion) { + c.bench_function("array comb c4", move |b| { + b.iter(|| { + for [i, j, k, l] in (0..N4).array_combinations() { + black_box(i + j + k + l); + } + }) + }); +} + + criterion_group!( benches, tuple_comb_for1, @@ -109,5 +154,10 @@ criterion_group!( tuple_comb_c2, tuple_comb_c3, tuple_comb_c4, + array_comb_c1, + array_comb_c2, + array_comb_c3, + array_comb_c4, ); + criterion_main!(benches); diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs index dfc68978f..f2bc137df 100644 --- a/src/adaptors/mod.rs +++ b/src/adaptors/mod.rs @@ -812,6 +812,72 @@ impl_tuple_combination!(Tuple10Combination Tuple9Combination; a b c d e f g h i) impl_tuple_combination!(Tuple11Combination Tuple10Combination; a b c d e f g h i j); impl_tuple_combination!(Tuple12Combination Tuple11Combination; a b c d e f g h i j k); +/// An iterator to iterate through all combinations in an iterator that produces arrays +/// of a specific size. +/// +/// See [`.array_combinations()`](crate::Itertools::array_combinations) for more +/// information. +#[derive(Debug, Clone)] +pub struct ArrayCombinations { + iter: I, + buf: Vec, + indices: [usize; R], +} + +impl ArrayCombinations { + /// Create a new `ArrayCombinations` from an iterator. + pub fn new(iter: I) -> Self { + let indices = array_init::array_init(|i| i); + let buf = Vec::new(); + + Self { iter, buf, indices } + } +} + +impl Iterator for ArrayCombinations +where + I::Item: Clone, +{ + type Item = [I::Item; R]; + + fn next(&mut self) -> Option { + if self.buf.is_empty() { + // If the buffer is empty, this is the first invocation of next + for _ in 0..R { + // If the source iter returns None, we won't have enough data + // for even 1 complete combination. So we can bail. + self.buf.push(self.iter.next()?); + } + } else if self.indices[0] + R == self.buf.len() { + // If the first index is as close to the end as possible + // eg: [0, 1, 2, 3, 4, 5] + // ^ ^ ^ + // then we can try get some more data. If there's no more data left + // then we've gone over all combinations of the underlying + // and we can bail + self.buf.push(self.iter.next()?); + + // Reset the indices + for i in 0..R - 1 { + self.indices[i] = i; + } + self.indices[R - 1] += 1; + } else { + let mut i = R - 2; + while i > 0 && self.indices[i] + R == self.buf.len() + i { + i -= 1; + } + + self.indices[i] += 1; + for j in i + 1..R-1 { + self.indices[j] = self.indices[j - 1] + 1; + } + } + + Some(array_init::array_init(|i| self.buf[self.indices[i]].clone())) + } +} + /// An iterator adapter to filter values within a nested `Result::Ok`. /// /// See [`.filter_ok()`](crate::Itertools::filter_ok) for more information. diff --git a/src/lib.rs b/src/lib.rs index 1a6632e20..8b9cbe268 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -105,6 +105,7 @@ pub mod structs { WhileSome, Coalesce, TupleCombinations, + ArrayCombinations, Positions, Update, }; @@ -1449,6 +1450,45 @@ 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` + /// + /// ``` + /// 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), (2, 3), (1, 4), (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::RangeFrom; + /// + /// let it: ArrayCombinations, 2> = (1..).array_combinations(); + /// itertools::assert_equal(it.take(6), vec![[1, 2], [1, 3], [2, 3], [1, 4], [2, 4], [3, 4]]); + /// ``` + fn array_combinations(self) -> ArrayCombinations + where Self: Sized, + Self::Item: Clone, + { + ArrayCombinations::new(self) + } + /// Return an iterator adaptor that iterates over the `k`-length combinations of /// the elements from an iterator. /// diff --git a/tests/quick.rs b/tests/quick.rs index 7769cb432..49d8fb6b6 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -895,6 +895,22 @@ quickcheck! { } } +quickcheck! { + fn equal_combinations_array(it: Iter) -> bool { + let values = it.clone().collect_vec(); + let mut cmb = it.array_combinations(); + for j in 1..values.len() { + for i in 0..j { + let pair = [values[i], values[j]]; + if pair != cmb.next().unwrap() { + return false; + } + } + } + cmb.next() == None + } +} + quickcheck! { fn size_pad_tail(it: Iter, pad: u8) -> bool { correct_size_hint(it.clone().pad_using(pad as usize, |_| 0)) &&