math/
lib.rs

1//! Precomputed modular binomial coefficients: `nCr mod p`.
2//!
3//! [`ModBinomial::new`] builds factorials and inverse factorials up to `max_n`
4//! in `O(max_n)` time using Fermat's little theorem for the modular inverse
5//! (`modulo` must be prime). Each subsequent [`ModBinomial::ncr`] call is `O(1)`.
6
7/// Binomial-coefficient table modulo a prime `modulo`.
8pub struct ModBinomial {
9    fact: Vec<i64>,
10    inv_fact: Vec<i64>,
11    modulo: i64,
12}
13
14impl ModBinomial {
15    /// Build the table for all `n` in `0..=max_n`. `modulo` must be prime
16    /// (the common choice is `1_000_000_007` or `998_244_353`).
17    pub fn new(max_n: usize, modulo: i64) -> Self {
18        let mut fact = vec![0; max_n + 1];
19        let mut inv_fact = vec![0; max_n + 1];
20
21        fact[0] = 1;
22        inv_fact[0] = 1;
23
24        for i in 1..=max_n {
25            fact[i] = (fact[i - 1] * i as i64) % modulo;
26        }
27
28        inv_fact[max_n] = Self::power(fact[max_n], modulo - 2, modulo);
29        for i in (1..=max_n).rev() {
30            inv_fact[i - 1] = (inv_fact[i] * i as i64) % modulo;
31        }
32
33        Self {
34            fact,
35            inv_fact,
36            modulo,
37        }
38    }
39
40    fn power(mut base: i64, mut exp: i64, modulo: i64) -> i64 {
41        let mut res = 1;
42        base %= modulo;
43        while exp > 0 {
44            if exp % 2 == 1 {
45                res = (res * base) % modulo;
46            }
47            base = (base * base) % modulo;
48            exp /= 2;
49        }
50        res
51    }
52
53    /// Compute `nCr mod modulo`. Returns `0` when `r > n`.
54    pub fn ncr(&self, n: usize, r: usize) -> i64 {
55        if r > n {
56            return 0;
57        }
58        let denom = (self.inv_fact[r] * self.inv_fact[n - r]) % self.modulo;
59        (self.fact[n] * denom) % self.modulo
60    }
61}
math/lib.rs

math/
lib.rs
