haar_lib/math/
stirling_second.rs

1//! 第二種スターリング数$S(n, 0), \dots, S(n, n)$を列挙する。
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/stirling_number_of_the_second_kind>
5
6use crate::math::ntt::NTT;
7use crate::num::const_modint::*;
8
9/// 第二種スターリング数$S(n, 0), \dots, S(n, n)$を列挙する。
10pub fn stirling_second<const P: u32, const PR: u32>(
11    n: usize,
12    ntt: &NTT<P, PR>,
13) -> Vec<ConstModInt<P>> {
14    let ff = ConstModIntBuilder;
15    let mut a = vec![ff.from_u64(0); n + 1];
16    let mut b = vec![ff.from_u64(0); n + 1];
17    let mut m = vec![0; n + 1];
18
19    for i in 2..=n {
20        if m[i] == 0 {
21            for j in (2 * i..=n).step_by(i) {
22                m[j] = i;
23            }
24        }
25    }
26
27    for i in 0..=n {
28        if m[i] == 0 {
29            a[i] = ff.from_u64(i as u64).pow(n as u64);
30        } else {
31            a[i] = a[m[i]] * a[i / m[i]];
32        }
33    }
34
35    let mut f = (1..=n)
36        .fold(ff.from_u64(1), |x, y| x * ff.from_u64(y as u64))
37        .inv();
38
39    for i in (0..=n).rev() {
40        a[i] *= f;
41        b[i] = f;
42        f *= ff.from_u64(i as u64);
43
44        if i % 2 == 1 {
45            b[i] = -b[i];
46        }
47    }
48
49    let mut ret = ntt.convolve(a, b);
50    ret.truncate(n + 1);
51    ret
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    use crate::math::{ntt::NTT998244353, stirling_second_table::stirling_second_table};
59
60    #[test]
61    fn test() {
62        let n = 100;
63        let ans = stirling_second_table(n, ConstModIntBuilder);
64
65        let ntt = NTT998244353::new();
66        for i in 0..=n {
67            assert_eq!(stirling_second(i, &ntt), ans[i][0..=i]);
68        }
69    }
70}
haar_lib/math/stirling_second.rs

haar_lib/math/
stirling_second.rs
