Auto merge of #272 - vks:binomial-coeffs, r=cuviper

Implement an iterator over the binomial coefficients

I'm not very happy with the excessive cloning, but to fix it the bounds on the type parameters would have to be excessive. We probably need something like [this](https://github.com/vks/discrete-log/blob/master/src/main.rs#L90) in `num-traits`.
This commit is contained in:
Homu 2017-03-31 10:29:14 +09:00
commit f63c933737
1 changed files with 121 additions and 1 deletions

View File

@ -667,14 +667,91 @@ impl_integer_for_usize!(u32, test_integer_u32);
impl_integer_for_usize!(u64, test_integer_u64);
impl_integer_for_usize!(usize, test_integer_usize);
/// An iterator over binomial coefficients.
pub struct IterBinomial<T> {
a: T,
n: T,
k: T,
}
impl<T> IterBinomial<T>
where T: Integer,
{
/// For a given n, iterate over all binomial coefficients binomial(n, k), for k=0...n.
///
/// Note that this might overflow, depending on `T`. For the primitive
/// integer types, the following n are the largest ones for which there will
/// be no overflow:
///
/// type | n
/// -----|---
/// u8 | 10
/// i8 | 9
/// u16 | 18
/// i16 | 17
/// u32 | 34
/// i32 | 33
/// u64 | 67
/// i64 | 66
///
/// For larger n, `T` should be a bigint type.
pub fn new(n: T) -> IterBinomial<T> {
IterBinomial {
k: T::zero(), a: T::one(), n: n
}
}
}
impl<T> Iterator for IterBinomial<T>
where T: Integer + Clone
{
type Item = T;
fn next(&mut self) -> Option<T> {
if self.k > self.n {
return None;
}
self.a = if !self.k.is_zero() {
multiply_and_divide(
self.a.clone(),
self.n.clone() - self.k.clone() + T::one(),
self.k.clone()
)
} else {
T::one()
};
self.k = self.k.clone() + T::one();
Some(self.a.clone())
}
}
/// Calculate r * a / b, avoiding overflows and fractions.
///
/// Assumes that b divides r * a evenly.
fn multiply_and_divide<T: Integer + Clone>(r: T, a: T, b: T) -> T {
// See http://blog.plover.com/math/choose-2.html for the idea.
let g = gcd(r.clone(), b.clone());
(r/g.clone() * a) / (b/g)
r/g.clone() * (a / (b/g))
}
/// Calculate the binomial coefficient.
///
/// Note that this might overflow, depending on `T`. For the primitive integer
/// types, the following n are the largest ones possible such that there will
/// be no overflow for any k:
///
/// type | n
/// -----|---
/// u8 | 10
/// i8 | 9
/// u16 | 18
/// i16 | 17
/// u32 | 34
/// i32 | 33
/// u64 | 67
/// i64 | 66
///
/// For larger n, consider using a bigint type for `T`.
pub fn binomial<T: Integer + Clone>(mut n: T, k: T) -> T {
// See http://blog.plover.com/math/choose.html for the idea.
if k > n {
@ -737,6 +814,49 @@ fn test_lcm_overflow() {
check!(u64, 0x8000_0000_0000_0000, 0x02, 0x8000_0000_0000_0000);
}
#[test]
fn test_iter_binomial() {
macro_rules! check_simple {
($t:ty) => { {
let n: $t = 3;
let c: Vec<_> = IterBinomial::new(n).collect();
let expected = vec![1, 3, 3, 1];
assert_eq!(c, expected);
} }
}
check_simple!(u8);
check_simple!(i8);
check_simple!(u16);
check_simple!(i16);
check_simple!(u32);
check_simple!(i32);
check_simple!(u64);
check_simple!(i64);
macro_rules! check_binomial {
($t:ty, $n:expr) => { {
let n: $t = $n;
let c: Vec<_> = IterBinomial::new(n).collect();
let mut k: $t = 0;
for b in c {
assert_eq!(b, binomial(n, k));
k += 1;
}
} }
}
// Check the largest n for which there is no overflow.
check_binomial!(u8, 10);
check_binomial!(i8, 9);
check_binomial!(u16, 18);
check_binomial!(i16, 17);
check_binomial!(u32, 34);
check_binomial!(i32, 33);
check_binomial!(u64, 67);
check_binomial!(i64, 66);
}
#[test]
fn test_binomial() {
macro_rules! check {