atcoder/lib/math/
combinatorics.rs

1//! 組合せ計算に関するアルゴリズムです。
2
3use super::mod_int::ModInt;
4
5/// 階乗と逆階乗を前計算し、順列・組合せを高速に計算します。
6///
7/// 法 `MOD` 上の階乗・順列・組合せを前計算して扱います。
8///
9/// `new(max_n)` は `0..=max_n` の階乗と逆階乗を保存します。以後、引数 `n` がこの範囲内なら
10/// `n!`、順列 `nPr`、組合せ `nCr` を `O(1)` で求められます。`r > n` の順列・組合せは `0` です。
11///
12/// `MOD` は素数であり、`max_n < MOD` を満たす必要があります。
13/// 構築の計算量は `O(max_n + log MOD)`、使用メモリは `O(max_n)` です。
14///
15/// # Panics
16///
17/// - `MOD <= 1` の場合
18/// - `max_n >= MOD` の場合
19/// - 前計算範囲を超える `n` を指定した場合
20///
21/// # Examples
22///
23/// ```
24/// use atcoder::math::combinatorics::Combinatorics;
25///
26/// let comb = Combinatorics::<998_244_353>::new(10);
27///
28/// assert_eq!(comb.factorial(5).value(), 120);
29/// assert_eq!(comb.permutation(5, 2).value(), 20);
30/// assert_eq!(comb.combination(5, 2).value(), 10);
31/// ```
32#[derive(Debug, Clone)]
33pub struct Combinatorics<const MOD: i64> {
34    factorial: Vec<ModInt<MOD>>,
35    inverse_factorial: Vec<ModInt<MOD>>,
36}
37
38impl<const MOD: i64> Combinatorics<MOD> {
39    /// `0..=max_n` の階乗と逆階乗を法 `MOD` で前計算します。
40    ///
41    /// `MOD` は `1` より大きい素数、かつ `max_n < MOD` でなければなりません。
42    ///
43    /// 計算量は `O(max_n + log MOD)`、使用メモリは `O(max_n)` です。
44    ///
45    /// # Panics
46    ///
47    /// `MOD <= 1`、`max_n >= MOD`、または `max_n` が `i64` に収まらない場合に panic します。
48    pub fn new(max_n: usize) -> Self {
49        assert!(MOD > 1, "modulus must be greater than one");
50
51        let max_n_i64 = i64::try_from(max_n).expect("max_n does not fit in i64");
52
53        assert!(max_n_i64 < MOD, "max_n must be smaller than modulus");
54
55        let mut factorial = vec![ModInt::new(1); max_n + 1];
56
57        for i in 1..=max_n {
58            factorial[i] = factorial[i - 1] * ModInt::new(i as i64);
59        }
60
61        // max_n! の逆元を求め、逆方向に掛け戻すことで
62        // すべての逆階乗を O(max_n) で計算する。
63        let mut inverse_factorial = vec![ModInt::new(1); max_n + 1];
64
65        inverse_factorial[max_n] = factorial[max_n].inv();
66
67        for i in (1..=max_n).rev() {
68            inverse_factorial[i - 1] = inverse_factorial[i] * ModInt::new(i as i64);
69        }
70
71        Self {
72            factorial,
73            inverse_factorial,
74        }
75    }
76
77    /// 前計算した最大の引数 `n` を返します。利用可能な範囲は `0..=max_n()` です。
78    pub fn max_n(&self) -> usize {
79        self.factorial.len() - 1
80    }
81
82    /// `n!` を法 `MOD` で返します。
83    ///
84    /// 計算量は `O(1)` です。
85    ///
86    /// # Panics
87    ///
88    /// `n > max_n()` の場合に panic します。
89    pub fn factorial(&self, n: usize) -> ModInt<MOD> {
90        self.assert_in_range(n);
91        self.factorial[n]
92    }
93
94    /// 順列 `nPr = n! / (n-r)!` を法 `MOD` で返します。
95    ///
96    /// `r > n` の場合は `0` を返します。`n` は前計算範囲内である必要があります。
97    /// 計算量は `O(1)` です。
98    pub fn permutation(&self, n: usize, r: usize) -> ModInt<MOD> {
99        self.assert_in_range(n);
100
101        if r > n {
102            return ModInt::new(0);
103        }
104
105        self.factorial[n] * self.inverse_factorial[n - r]
106    }
107
108    /// 組合せ `nCr` を法 `MOD` で返します。
109    ///
110    /// `r > n` の場合は `0` を返します。`n` は前計算範囲内である必要があります。
111    /// 計算量は `O(1)` です。
112    pub fn combination(&self, n: usize, r: usize) -> ModInt<MOD> {
113        self.assert_in_range(n);
114
115        if r > n {
116            return ModInt::new(0);
117        }
118
119        self.factorial[n] * self.inverse_factorial[r] * self.inverse_factorial[n - r]
120    }
121
122    fn assert_in_range(&self, n: usize) {
123        assert!(
124            n <= self.max_n(),
125            "n exceeds the precomputed range: n = {n}, max_n = {}",
126            self.max_n(),
127        );
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::Combinatorics;
134
135    type Comb = Combinatorics<998_244_353>;
136
137    #[test]
138    fn calculates_factorials_and_combinations() {
139        let comb = Comb::new(10);
140
141        assert_eq!(comb.factorial(0).value(), 1);
142        assert_eq!(comb.factorial(5).value(), 120);
143
144        assert_eq!(comb.permutation(5, 0).value(), 1);
145        assert_eq!(comb.permutation(5, 2).value(), 20);
146        assert_eq!(comb.permutation(5, 6).value(), 0);
147
148        assert_eq!(comb.combination(5, 0).value(), 1);
149        assert_eq!(comb.combination(5, 2).value(), 10);
150        assert_eq!(comb.combination(5, 5).value(), 1);
151        assert_eq!(comb.combination(5, 6).value(), 0);
152    }
153
154    #[test]
155    fn satisfies_pascals_identity() {
156        let comb = Comb::new(30);
157
158        for n in 1..=30 {
159            for r in 1..n {
160                let lhs = comb.combination(n, r).value();
161
162                let rhs = (comb.combination(n - 1, r - 1).value()
163                    + comb.combination(n - 1, r).value())
164                    % 998_244_353;
165
166                assert_eq!(lhs, rhs, "n = {n}, r = {r}");
167            }
168        }
169    }
170}
atcoder/lib/math/combinatorics.rs

atcoder/lib/math/
combinatorics.rs
