cp_library_rs/number_theory/
comb.rs

1//! 階乗を前計算する（Modint構造体に依存）
2
3use std::ops::{Add, Mul, Neg, Sub};
4
5use crate::number_theory::modint::{M107, M998, MOD107, MOD998};
6
7use num_traits::{One, Zero};
8
9/// 確率の値となりうる数
10pub trait NumProbability<Rhs = Self, Output = Self>:
11    Clone
12    + Copy
13    + From<usize>
14    + Neg<Output = Output>
15    + Add<Rhs, Output = Output>
16    + Sub<Rhs, Output = Output>
17    + Mul<Rhs, Output = Output>
18    + Mul<usize, Output = Output>
19    + Zero
20    + One
21    + PartialEq
22{
23    const MOD: usize;
24}
25impl NumProbability for M107 {
26    const MOD: usize = MOD107;
27}
28impl NumProbability for M998 {
29    const MOD: usize = MOD998;
30}
31
32/// 二項係数を高速に求める
33/// - 前計算:  $`O(N)`$ 時間
34/// - クエリ:  $`O(1)`$ 時間
35pub struct Comb<N: NumProbability> {
36    fac: Vec<N>,
37    finv: Vec<N>,
38}
39
40impl<N: NumProbability> Comb<N> {
41    /// サイズ`max_size`で配列を初期化する
42    pub fn new(max_size: usize) -> Self {
43        let mut fac = vec![N::one(); max_size];
44        let mut finv = vec![N::one(); max_size];
45        let mut inv = vec![N::one(); max_size];
46        for i in 2..max_size {
47            fac[i] = fac[i - 1] * i;
48            inv[i] = -N::from(N::MOD / i) * inv[N::MOD % i];
49            finv[i] = finv[i - 1] * inv[i];
50        }
51        Comb { fac, finv }
52    }
53
54    /// n 個から r 個選ぶ組合せ
55    ///
56    /// - 時間計算量: $`O(1)`$
57    pub fn comb(&self, n: usize, r: usize) -> N {
58        if n < r {
59            return 0.into();
60        }
61        self.fac[n] * self.finv[r] * self.finv[n - r]
62    }
63
64    /// n 個から r 個を選び並べる順列
65    ///
66    /// - 時間計算量: $`O(1)`$
67    pub fn perm(&self, n: usize, r: usize) -> N {
68        if n < r {
69            return 0.into();
70        }
71        self.fac[n] * self.finv[n - r]
72    }
73
74    /// 重複組合せ
75    ///
76    /// - balls 個の区別しない玉を boxes 個の区別する箱に入れる組合せ
77    /// - 時間計算量: $`O(1)`$
78    pub fn comb_with_rep(&self, balls: usize, boxes: usize) -> N {
79        self.comb(balls + boxes - 1, balls)
80    }
81}
cp_library_rs/number_theory/comb.rs

cp_library_rs/number_theory/
comb.rs
