Implement multinomial

This commit is contained in:
Vinzent Steinberg 2017-03-17 17:25:26 +01:00
parent 2adf018aa6
commit ebe3f85976
1 changed files with 91 additions and 0 deletions

View File

@ -16,6 +16,8 @@
extern crate num_traits as traits; extern crate num_traits as traits;
use std::ops::Add;
use traits::{Num, Signed}; use traits::{Num, Signed};
pub trait Integer: Sized + Num + PartialOrd + Ord + Eq { pub trait Integer: Sized + Num + PartialOrd + Ord + Eq {
@ -694,6 +696,22 @@ pub fn binomial<T: Integer + Clone>(mut n: T, k: T) -> T {
r r
} }
/// Calculate the multinomial coefficient.
///
/// Panics if no integers are specified.
pub fn multinomial<T: Integer + Clone>(k: &[T]) -> T
where for<'a> T: Add<&'a T, Output = T>
{
assert!(k.len() > 0, "need at least one integer");
let mut r = T::one();
let mut p = T::zero();
for i in k {
p = p + i;
r = r * binomial(p.clone(), i.clone());
}
r
}
#[test] #[test]
fn test_lcm_overflow() { fn test_lcm_overflow() {
macro_rules! check { macro_rules! check {
@ -777,3 +795,76 @@ fn test_binomial() {
check!(i64, 0, 0, 1); check!(i64, 0, 0, 1);
check!(i64, 2, 3, 0); check!(i64, 2, 3, 0);
} }
#[test]
fn test_multinomial() {
macro_rules! check_binomial {
($t:ty, $k:expr) => { {
let n: $t = $k.iter().sum();
let k: &[$t] = $k;
assert_eq!(k.len(), 2);
assert_eq!(multinomial(k), binomial(n, k[0]));
} }
}
check_binomial!(u8, &[4, 5]);
check_binomial!(i8, &[4, 5]);
check_binomial!(u16, &[2, 98]);
check_binomial!(u16, &[4, 10]);
check_binomial!(i16, &[2, 98]);
check_binomial!(i16, &[4, 10]);
check_binomial!(u32, &[2, 98]);
check_binomial!(u32, &[11, 24]);
check_binomial!(u32, &[4, 10]);
check_binomial!(i32, &[2, 98]);
check_binomial!(i32, &[11, 24]);
check_binomial!(i32, &[4, 10]);
check_binomial!(u64, &[2, 98]);
check_binomial!(u64, &[11, 24]);
check_binomial!(u64, &[4, 10]);
check_binomial!(i64, &[2, 98]);
check_binomial!(i64, &[11, 24]);
check_binomial!(i64, &[4, 10]);
macro_rules! check_multinomial {
($t:ty, $k:expr, $r:expr) => { {
let k: &[$t] = $k;
let expected: $t = $r;
assert_eq!(multinomial(k), expected);
} }
}
check_multinomial!(u8, &[2, 1, 2], 30);
check_multinomial!(u8, &[2, 3, 0], 10);
check_multinomial!(i8, &[2, 1, 2], 30);
check_multinomial!(i8, &[2, 3, 0], 10);
check_multinomial!(u16, &[2, 1, 2], 30);
check_multinomial!(u16, &[2, 3, 0], 10);
check_multinomial!(i16, &[2, 1, 2], 30);
check_multinomial!(i16, &[2, 3, 0], 10);
check_multinomial!(u32, &[2, 1, 2], 30);
check_multinomial!(u32, &[2, 3, 0], 10);
check_multinomial!(i32, &[2, 1, 2], 30);
check_multinomial!(i32, &[2, 3, 0], 10);
check_multinomial!(u64, &[2, 1, 2], 30);
check_multinomial!(u64, &[2, 3, 0], 10);
check_multinomial!(i64, &[2, 1, 2], 30);
check_multinomial!(i64, &[2, 3, 0], 10);
check_multinomial!(u64, &[0], 1);
check_multinomial!(u64, &[12345], 1);
}