haar_lib/math/combinatorics/
stirling_second_small_p.rs

1//! 第二種スターリング数$S(n, k)$を計算する。
2//!
3//! # References
4//! - <https://maspypy.com/stirling-%e6%95%b0%e3%82%92-p-%e3%81%a7%e5%89%b2%e3%81%a3%e3%81%9f%e4%bd%99%e3%82%8a%e3%81%ae%e8%a8%88%e7%ae%97>
5//! # Problems
6//! - <https://judge.yosupo.jp/problem/stirling_number_of_the_second_kind_small_p_large_n>
7use crate::math::combinatorics::binomial_coefficient::ExtLucas;
8
9/// 第二種スターリング数$S(n, k)$を計算する。
10pub struct StirlingSecondSmallP {
11    p: u64,
12    bc: ExtLucas,
13    s: Vec<Vec<u64>>,
14}
15
16impl StirlingSecondSmallP {
17    /// $\mod p$ ($p$は素数)で前計算をする。
18    ///
19    /// **Time complexity** $O(p^2)$
20    ///
21    /// **Space complexity** $O(p^2)$
22    pub fn new(p: u64) -> Self {
23        let bc = ExtLucas::new(p, 1);
24        let mut s = vec![vec![0; p as usize + 1]; p as usize + 1];
25
26        s[0][0] = 1;
27
28        for (i, si) in s.iter_mut().enumerate().skip(1) {
29            si[1] = 1;
30            si[i] = 1;
31        }
32
33        for i in 3..=p as usize {
34            for j in 2..i {
35                s[i][j] = (s[i - 1][j - 1] + j as u64 * s[i - 1][j] % p) % p;
36            }
37        }
38
39        Self { p, bc, s }
40    }
41
42    /// $S(n, k)$を計算する。
43    ///
44    /// **Time complexity** $O(1)$
45    pub fn calc(&self, n: u64, k: u64) -> u64 {
46        if n <= self.p && k <= self.p {
47            return self.s[n as usize][k as usize];
48        }
49
50        let i = k / self.p;
51        let j = k % self.p;
52
53        let mut b = (n - i) % (self.p - 1);
54        if b == 0 {
55            b = self.p - 1;
56        }
57        let a = (n - i - b) / (self.p - 1);
58
59        if b == self.p - 1 {
60            self.bc.calc(a, i) * self.s[self.p as usize - 1][j as usize] % self.p
61                + self.bc.calc(a, i - 1) * self.s[0][j as usize] % self.p
62        } else {
63            self.bc.calc(a, i) * self.s[b as usize][j as usize] % self.p
64        }
65    }
66}
haar_lib/math/combinatorics/stirling_second_small_p.rs

haar_lib/math/combinatorics/
stirling_second_small_p.rs
