diff --git a/src/function/factorial.rs b/src/function/factorial.rs index d8338884..52c2de1a 100644 --- a/src/function/factorial.rs +++ b/src/function/factorial.rs @@ -40,13 +40,43 @@ pub fn ln_factorial(x: u64) -> f64 { /// /// # Remarks /// -/// Returns `0.0` if `k > n` +/// Returns `0.0` if `k > n`. +/// +/// The result is exact whenever `C(n, k)` fits in a `u128`, and correctly +/// rounded to the nearest `f64` beyond `2^53`. Larger coefficients fall back +/// to `exp(ln_factorial(n) - ln_factorial(k) - ln_factorial(n - k))`, which +/// is accurate to about `1e-12` relative. pub fn binomial(n: u64, k: u64) -> f64 { if k > n { - 0.0 - } else { - (0.5 + (ln_factorial(n) - ln_factorial(k) - ln_factorial(n - k)).exp()).floor() + return 0.0; + } + match binomial_u128(n, k) { + Some(exact) => exact as f64, + // Overflow: the result exceeds ~3.4e38 / n, far past 2^53, so + // exactness is unattainable anyway; fall back to logs. + None => (ln_factorial(n) - ln_factorial(k) - ln_factorial(n - k)) + .exp() + .round(), + } +} + +/// Computes `C(n, k)` exactly in integer arithmetic, or `None` if an +/// intermediate value overflows a `u128`. Requires `k <= n`. +/// +/// `C(n, i + 1) = C(n, i) * (n - i) / (i + 1)`, and the division is exact at +/// every step, so the whole computation stays in integer arithmetic. The +/// previous `f64` implementation rounded `exp(ln n! - ln k! - ln (n-k)!)` to +/// the nearest integer, which was off by up to ~1e5 (57 ulp) for coefficients +/// near 2^63 that f64 could still represent to full precision. +fn binomial_u128(n: u64, k: u64) -> Option { + // C(n, k) == C(n, n - k); walking up the smaller side minimises both the + // iteration count and the chance of overflowing the integer path. + let k = k.min(n - k); + let mut acc: u128 = 1; + for i in 0..k { + acc = acc.checked_mul((n - i) as u128)? / (i as u128 + 1); } + Some(acc) } /// Computes the natural logarithm of the binomial coefficient @@ -75,16 +105,32 @@ pub fn multinomial(n: u64, ni: &[u64]) -> f64 { /// Computes the multinomial coefficient: `n choose n1, n2, n3, ...` /// /// Returns `None` if the elements in `ni` do not sum to `n`. +/// +/// # Remarks +/// +/// The result is exact whenever the coefficient (and every intermediate +/// prefix product) fits in a `u128`, and correctly rounded to the nearest +/// `f64` beyond `2^53`. Larger coefficients fall back to +/// `exp(ln n! - sum ln ni!)`, accurate to about `1e-12` relative. pub fn checked_multinomial(n: u64, ni: &[u64]) -> Option { - let (sum, ret) = ni.iter().fold((0, ln_factorial(n)), |acc, &x| { - (acc.0 + x, acc.1 - ln_factorial(x)) - }); + if ni.iter().sum::() != n { + return None; + } - if sum == n { - Some((0.5 + ret.exp()).floor()) - } else { - None + // n! / (n1! n2! ... nk!) == prod_i C(s_i, n_i) with s_i = n_1 + ... + n_i: + // a product of binomial coefficients, each computed exactly. + let mut acc: u128 = 1; + let mut prefix: u64 = 0; + for &k in ni { + prefix += k; + let Some(product) = binomial_u128(prefix, k).and_then(|c| acc.checked_mul(c)) else { + // Overflow: fall back to logs (the old implementation's only path). + let ret = ni.iter().fold(ln_factorial(n), |a, &x| a - ln_factorial(x)); + return Some(ret.exp().round()); + }; + acc = product; } + Some(acc as f64) } // Initialization for pre-computed cache of 171 factorial @@ -158,6 +204,49 @@ mod tests { assert_eq!(binomial(5, 7), 0.0); } + /// Every `C(n, k)` that fits in a `u128` must round-trip exactly (the + /// conversion `u128 -> f64` is correctly rounded, so `expected as f64` is + /// the best possible double). The old `exp(ln ...)`-based implementation + /// failed this for e.g. `C(50, 25)` (off by 2) and `C(67, 33)` (off by + /// 116736). + #[test] + fn test_binomial_is_exact_where_representable() { + for n in 0..=170u64 { + let mut expected: u128 = 1; + for k in 0..=n / 2 { + assert_eq!( + binomial(n, k), + expected as f64, + "C({n}, {k}) should be {expected}" + ); + assert_eq!(binomial(n, n - k), expected as f64, "C({n}, {}) symmetry", n - k); + let Some(product) = expected.checked_mul((n - k) as u128) else { + break; + }; + expected = product / (k as u128 + 1); + } + } + } + + /// Coefficients too large for the integer path fall back to logs; check the + /// fallback is close and consistent with `ln_binomial`. + #[test] + fn test_binomial_log_fallback() { + // C(200, 100) = 9.0548514656103281165404177077e58 (overflows u128) + prec::assert_relative_eq!( + binomial(200, 100), + 9.0548514656103281165e58, + epsilon = 0.0, + max_relative = 1e-11 + ); + prec::assert_relative_eq!( + binomial(1000, 500), + ln_binomial(1000, 500).exp(), + epsilon = 0.0, + max_relative = 1e-11 + ); + } + #[test] fn test_ln_binomial() { assert_eq!(ln_binomial(1, 1), 1f64.ln()); @@ -176,6 +265,36 @@ mod tests { assert_eq!(35.0, multinomial(7, &[3, 4])); } + /// A two-part multinomial is a binomial coefficient; a three-part one has + /// the closed form `C(n, a) * C(n - a, b)`. Both must be exact where + /// representable (mirrors `test_binomial_is_exact_where_representable`). + #[test] + fn test_multinomial_is_exact_where_representable() { + for n in 0..=170u64 { + for k in 0..=n / 2 { + assert_eq!(multinomial(n, &[k, n - k]), binomial(n, k), "n={n} k={k}"); + } + } + // 60! / (20!)^3 = 577831214478475823831865900 (fits u128); the + // conversion `u128 -> f64` is correctly rounded: + assert_eq!( + multinomial(60, &[20, 20, 20]), + 577831214478475823831865900_u128 as f64 + ); + } + + #[test] + fn test_multinomial_log_fallback() { + // 300! / (100!)^3 overflows u128; check against ln-space value + let ln_ref = ln_factorial(300) - 3.0 * ln_factorial(100); + prec::assert_relative_eq!( + multinomial(300, &[100, 100, 100]), + ln_ref.exp(), + epsilon = 0.0, + max_relative = 1e-11 + ); + } + #[test] #[should_panic] fn test_multinomial_bad_ni() {