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>(n: usize) -> Vec<ConstModInt<P>> {
11    let ntt = NTT::<P, PR>::new();
12    let ff = ConstModIntBuilder;
13    let mut a = vec![ff.from_u64(0); n + 1];
14    let mut b = vec![ff.from_u64(0); n + 1];
15    let mut m = vec![0; n + 1];
16
17    for i in 2..=n {
18        if m[i] == 0 {
19            for j in (2 * i..=n).step_by(i) {
20                m[j] = i;
21            }
22        }
23    }
24
25    for i in 0..=n {
26        if m[i] == 0 {
27            a[i] = ff.from_u64(i as u64).pow(n as u64);
28        } else {
29            a[i] = a[m[i]] * a[i / m[i]];
30        }
31    }
32
33    let mut f = (1..=n)
34        .fold(ff.from_u64(1), |x, y| x * ff.from_u64(y as u64))
35        .inv();
36
37    for i in (0..=n).rev() {
38        a[i] *= f;
39        b[i] = f;
40        f *= ff.from_u64(i as u64);
41
42        if i % 2 == 1 {
43            b[i] = -b[i];
44        }
45    }
46
47    let mut ret = ntt.convolve(a, b);
48    ret.truncate(n + 1);
49    ret
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55
56    use crate::math::stirling_second_table::stirling_second_table;
57
58    #[test]
59    fn test() {
60        let n = 100;
61        let ans = stirling_second_table(n, ConstModIntBuilder);
62
63        for i in 0..=n {
64            assert_eq!(stirling_second::<998244353, 3>(i), ans[i][0..=i]);
65        }
66    }
67}
haar_lib/math/stirling_second.rs

haar_lib/math/
stirling_second.rs
