diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs index 4bc50eaa0..90cf46eed 100644 --- a/src/adaptors/mod.rs +++ b/src/adaptors/mod.rs @@ -14,7 +14,7 @@ pub use self::map::{map_into, map_ok, MapInto, MapOk}; #[cfg(feature = "use_alloc")] pub use self::multi_product::*; -use crate::size_hint; +use crate::size_hint::{self, SizeHint}; use std::fmt; use std::iter::{FromIterator, Fuse, FusedIterator}; use std::marker::PhantomData; @@ -626,6 +626,14 @@ where fn next(&mut self) -> Option { self.iter.next() } + + fn size_hint(&self) -> SizeHint { + self.iter.size_hint() + } + + fn count(self) -> usize { + self.iter.count() + } } impl FusedIterator for TupleCombinations @@ -652,6 +660,14 @@ impl Iterator for Tuple1Combination { fn next(&mut self) -> Option { self.iter.next().map(|x| (x,)) } + + fn size_hint(&self) -> SizeHint { + self.iter.size_hint() + } + + fn count(self) -> usize { + self.iter.count() + } } impl HasCombination for (I::Item,) { @@ -701,6 +717,20 @@ macro_rules! impl_tuple_combination { }) } } + + fn size_hint(&self) -> SizeHint { + const K: usize = 1 + count_ident!($($X,)*); + let (mut n_min, mut n_max) = self.iter.size_hint(); + n_min = checked_binomial(n_min, K).unwrap_or(usize::MAX); + n_max = n_max.and_then(|n| checked_binomial(n, K)); + size_hint::add(self.c.size_hint(), (n_min, n_max)) + } + + fn count(self) -> usize { + const K: usize = 1 + count_ident!($($X,)*); + let n = self.iter.count(); + checked_binomial(n, K).unwrap() + self.c.count() + } } impl HasCombination for (A, $(ignore_ident!($X, A)),*) @@ -736,6 +766,42 @@ 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); +// https://en.wikipedia.org/wiki/Binomial_coefficient#In_programming_languages +pub(crate) fn checked_binomial(mut n: usize, mut k: usize) -> Option { + if n < k { + return Some(0); + } + // `factorial(n) / factorial(n - k) / factorial(k)` but trying to avoid it overflows: + k = (n - k).min(k); // symmetry + let mut c = 1; + for i in 1..=k { + c = (c / i) + .checked_mul(n)? + .checked_add((c % i).checked_mul(n)? / i)?; + n -= 1; + } + Some(c) +} + +#[test] +fn test_checked_binomial() { + // With the first row: [1, 0, 0, ...] and the first column full of 1s, we check + // row by row the recurrence relation of binomials (which is an equivalent definition). + // For n >= 1 and k >= 1 we have: + // binomial(n, k) == binomial(n - 1, k - 1) + binomial(n - 1, k) + const LIMIT: usize = 500; + let mut row = vec![Some(0); LIMIT + 1]; + row[0] = Some(1); + for n in 0..=LIMIT { + for k in 0..=LIMIT { + assert_eq!(row[k], checked_binomial(n, k)); + } + row = std::iter::once(Some(1)) + .chain((1..=LIMIT).map(|k| row[k - 1]?.checked_add(row[k]?))) + .collect(); + } +} + /// 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/combinations.rs b/src/combinations.rs index 332bcb3e2..1c834a74e 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -4,6 +4,8 @@ use std::iter::FusedIterator; use super::lazy_buffer::LazyBuffer; use alloc::vec::Vec; +use crate::adaptors::checked_binomial; + /// An iterator to iterate through all the `k`-length combinations in an iterator. /// /// See [`.combinations()`](crate::Itertools::combinations) for more information. @@ -160,42 +162,6 @@ where { } -// https://en.wikipedia.org/wiki/Binomial_coefficient#In_programming_languages -pub(crate) fn checked_binomial(mut n: usize, mut k: usize) -> Option { - if n < k { - return Some(0); - } - // `factorial(n) / factorial(n - k) / factorial(k)` but trying to avoid it overflows: - k = (n - k).min(k); // symmetry - let mut c = 1; - for i in 1..=k { - c = (c / i) - .checked_mul(n)? - .checked_add((c % i).checked_mul(n)? / i)?; - n -= 1; - } - Some(c) -} - -#[test] -fn test_checked_binomial() { - // With the first row: [1, 0, 0, ...] and the first column full of 1s, we check - // row by row the recurrence relation of binomials (which is an equivalent definition). - // For n >= 1 and k >= 1 we have: - // binomial(n, k) == binomial(n - 1, k - 1) + binomial(n - 1, k) - const LIMIT: usize = 500; - let mut row = vec![Some(0); LIMIT + 1]; - row[0] = Some(1); - for n in 0..=LIMIT { - for k in 0..=LIMIT { - assert_eq!(row[k], checked_binomial(n, k)); - } - row = std::iter::once(Some(1)) - .chain((1..=LIMIT).map(|k| row[k - 1]?.checked_add(row[k]?))) - .collect(); - } -} - /// 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 { let k = indices.len(); diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs index e4f1cda13..f1a0bc07d 100644 --- a/src/combinations_with_replacement.rs +++ b/src/combinations_with_replacement.rs @@ -3,7 +3,7 @@ use std::fmt; use std::iter::FusedIterator; use super::lazy_buffer::LazyBuffer; -use crate::combinations::checked_binomial; +use crate::adaptors::checked_binomial; /// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement. /// diff --git a/src/impl_macros.rs b/src/impl_macros.rs index 5fd78cb4d..7c0b0c22e 100644 --- a/src/impl_macros.rs +++ b/src/impl_macros.rs @@ -27,3 +27,8 @@ macro_rules! clone_fields { macro_rules! ignore_ident{ ($id:ident, $($t:tt)*) => {$($t)*}; } + +macro_rules! count_ident { + () => {0}; + ($i0:ident, $($i:ident,)*) => {1 + count_ident!($($i,)*)}; +} diff --git a/src/powerset.rs b/src/powerset.rs index 811cfce7f..a810f9582 100644 --- a/src/powerset.rs +++ b/src/powerset.rs @@ -3,7 +3,8 @@ use std::fmt; use std::iter::FusedIterator; use std::usize; -use super::combinations::{checked_binomial, combinations, Combinations}; +use super::combinations::{combinations, Combinations}; +use crate::adaptors::checked_binomial; use crate::size_hint::{self, SizeHint}; /// An iterator to iterate through the powerset of the elements from an iterator. diff --git a/src/tuple_impl.rs b/src/tuple_impl.rs index ba1c69df6..536f796fb 100644 --- a/src/tuple_impl.rs +++ b/src/tuple_impl.rs @@ -297,10 +297,6 @@ pub trait TupleCollect: Sized { fn left_shift_push(&mut self, item: Self::Item); } -macro_rules! count_ident{ - () => {0}; - ($i0:ident, $($i:ident,)*) => {1 + count_ident!($($i,)*)}; -} macro_rules! rev_for_each_ident{ ($m:ident, ) => {}; ($m:ident, $i0:ident, $($i:ident,)*) => { diff --git a/tests/quick.rs b/tests/quick.rs index 3dc23d704..5eec82844 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -901,8 +901,31 @@ quickcheck! { } quickcheck! { - fn size_combinations(it: Iter) -> bool { - correct_size_hint(it.tuple_combinations::<(_, _)>()) + fn size_combinations(a: Iter) -> bool { + let it = a.clone().tuple_combinations::<(_, _)>(); + correct_size_hint(it.clone()) && it.count() == binomial(a.count(), 2) + } + + fn exact_size_combinations_1(a: Vec) -> bool { + let it = a.iter().tuple_combinations::<(_,)>(); + exact_size_for_this(it.clone()) && it.count() == binomial(a.len(), 1) + } + fn exact_size_combinations_2(a: Vec) -> bool { + let it = a.iter().tuple_combinations::<(_, _)>(); + exact_size_for_this(it.clone()) && it.count() == binomial(a.len(), 2) + } + fn exact_size_combinations_3(mut a: Vec) -> bool { + a.truncate(15); + let it = a.iter().tuple_combinations::<(_, _, _)>(); + exact_size_for_this(it.clone()) && it.count() == binomial(a.len(), 3) + } +} + +fn binomial(n: usize, k: usize) -> usize { + if k > n { + 0 + } else { + (n - k + 1..=n).product::() / (1..=k).product::() } } diff --git a/tests/specializations.rs b/tests/specializations.rs index efc8ba6ef..57a85e5cc 100644 --- a/tests/specializations.rs +++ b/tests/specializations.rs @@ -72,6 +72,12 @@ where } quickcheck! { + fn tuple_combinations(v: Vec) -> () { + let mut v = v; + v.truncate(10); + test_specializations(&v.iter().tuple_combinations::<(_, _, _)>()); + } + fn intersperse(v: Vec) -> () { test_specializations(&v.into_iter().intersperse(0)); }