Skip to main content

nekolib/math/
mod_factorial_binom.rs

1//! 法 $p$ での二項係数。
2
3use super::mod_recip_table_;
4
5use mod_recip_table_::mod_recip_table_prime;
6
7/// 法 $p$ での二項係数。
8///
9/// # Examples
10/// ```
11/// use nekolib::math::ModFactorialBinom;
12///
13/// const MOD: u64 = 998244353;
14/// let mfb = ModFactorialBinom::new(10, MOD);
15///
16/// assert_eq!(mfb.factorial(5), 120);
17/// assert_eq!(mfb.factorial_recip(5), 856826403);
18/// assert_eq!(mfb.recip(3), 332748118);
19///
20/// assert_eq!(mfb.perm(6, 4), 360);
21/// assert_eq!(mfb.binom(7, 3), 35);
22/// ```
23///
24/// ```should_panic
25/// use nekolib::math::ModFactorialBinom;
26///
27/// const MOD: u64 = 998244353;
28/// let mfb = ModFactorialBinom::new(10, MOD);
29///
30/// mfb.recip(0);
31/// ```
32pub struct ModFactorialBinom {
33    p: u64,
34    fact: Vec<u64>,
35    fact_recip: Vec<u64>,
36}
37
38impl ModFactorialBinom {
39    /// $(0!, 1!, \\dots, n!)$ と $(0!^{-1}, 1!^{-1}, \\dots, n!^{-1})$ を
40    /// $p$ を法として計算する。
41    pub fn new(n: usize, p: u64) -> Self {
42        let mut fact = vec![1; n + 1];
43        for i in 2..=n {
44            fact[i] = fact[i - 1] * i as u64 % p;
45        }
46        let recip = mod_recip_table_prime(n as u64, p);
47        let mut fact_recip = vec![1; n + 1];
48        for i in 2..=n {
49            fact_recip[i] = fact_recip[i - 1] * recip[i] % p;
50        }
51        Self { p, fact, fact_recip }
52    }
53
54    /// $i! \\bmod p$ を返す。
55    pub fn factorial(&self, i: usize) -> u64 { self.fact[i] }
56
57    /// $i!^{-1} \\bmod p$ を返す。
58    pub fn factorial_recip(&self, i: usize) -> u64 { self.fact_recip[i] }
59
60    /// $i^{-1} \\bmod p$ を返す。
61    pub fn recip(&self, i: usize) -> u64 {
62        self.fact_recip[i] * self.fact[i - 1] % self.p
63    }
64
65    /// $i!/(i-j)! \\bmod p$ を返す。
66    pub fn perm(&self, i: usize, j: usize) -> u64 {
67        if i < j {
68            0
69        } else {
70            self.fact[i] * self.fact_recip[i - j] % self.p
71        }
72    }
73
74    /// $i!/(j!\\cdot (i-j)!) \\bmod p$ を返す。
75    pub fn binom(&self, i: usize, j: usize) -> u64 {
76        if i < j {
77            0
78        } else {
79            self.perm(i, j) * self.fact_recip[j] % self.p
80        }
81    }
82}