haar_lib/math/combinatorics/
stirling_first_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_first_kind_small_p_large_n>
7use crate::math::combinatorics::binomial_coefficient::ExtLucas;
8
9/// 符号付き第一種スターリング数$s(n, k)$を計算する。
10pub struct StirlingFirstSmallP {
11    p: u64,
12    bc: ExtLucas,
13    s: Vec<Vec<u64>>,
14}
15
16impl StirlingFirstSmallP {
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        s[0][0] = 1;
26
27        for i in 1..=p as usize {
28            for j in 1..=i {
29                let mut t = (i as u64 - 1) * s[i - 1][j] % p;
30                if t != 0 {
31                    t = p - t;
32                }
33
34                s[i][j] = t + s[i - 1][j - 1];
35                if s[i][j] >= p {
36                    s[i][j] %= p;
37                }
38            }
39        }
40
41        Self { p, bc, s }
42    }
43
44    /// $s(n, k)$を計算する。
45    ///
46    /// **Time complexity** $O(1)$
47    pub fn calc(&self, n: u64, k: u64) -> u64 {
48        let i = n / self.p;
49        let j = n % self.p;
50
51        let mut b = (k - i) % (self.p - 1);
52        if b == 0 && j > 0 {
53            b = self.p - 1;
54        }
55        let a = ((k - i) - b) / (self.p - 1);
56
57        let mut ret = self.bc.calc(i, a) * self.s[j as usize][b as usize] % self.p;
58        if (i - a) % 2 == 1 && ret != 0 {
59            ret = self.p - ret;
60        }
61
62        ret
63    }
64}
haar_lib/math/combinatorics/stirling_first_small_p.rs

haar_lib/math/combinatorics/
stirling_first_small_p.rs
