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:
commit
f63c933737
|
@ -667,14 +667,91 @@ impl_integer_for_usize!(u32, test_integer_u32);
|
||||||
impl_integer_for_usize!(u64, test_integer_u64);
|
impl_integer_for_usize!(u64, test_integer_u64);
|
||||||
impl_integer_for_usize!(usize, test_integer_usize);
|
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.
|
/// 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 {
|
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.
|
// See http://blog.plover.com/math/choose-2.html for the idea.
|
||||||
let g = gcd(r.clone(), b.clone());
|
let g = gcd(r.clone(), b.clone());
|
||||||
(r/g.clone() * a) / (b/g)
|
r/g.clone() * (a / (b/g))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Calculate the binomial coefficient.
|
/// 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 {
|
pub fn binomial<T: Integer + Clone>(mut n: T, k: T) -> T {
|
||||||
// See http://blog.plover.com/math/choose.html for the idea.
|
// See http://blog.plover.com/math/choose.html for the idea.
|
||||||
if k > n {
|
if k > n {
|
||||||
|
@ -737,6 +814,49 @@ fn test_lcm_overflow() {
|
||||||
check!(u64, 0x8000_0000_0000_0000, 0x02, 0x8000_0000_0000_0000);
|
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]
|
#[test]
|
||||||
fn test_binomial() {
|
fn test_binomial() {
|
||||||
macro_rules! check {
|
macro_rules! check {
|
||||||
|
|
Loading…
Reference in New Issue