haar_lib/math/combinatorics/
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::convolution::ntt::NTT;
7use crate::math::prime_mod::PrimeMod;
8use crate::num::const_modint::*;
9
10/// 第二種スターリング数$S(n, 0), \dots, S(n, n)$を列挙する。
11pub fn stirling_second<P: PrimeMod>(n: usize) -> Vec<ConstModInt<P>> {
12    let ntt = NTT::<P>::new();
13    let ff = ConstModIntBuilder::new();
14    let mut a = vec![ff.from_u64(0); n + 1];
15    let mut b = vec![ff.from_u64(0); n + 1];
16    let mut m = vec![0; n + 1];
17
18    for i in 2..=n {
19        if m[i] == 0 {
20            for j in (2 * i..=n).step_by(i) {
21                m[j] = i;
22            }
23        }
24    }
25
26    for i in 0..=n {
27        if m[i] == 0 {
28            a[i] = ff.from_u64(i as u64).pow(n as u64);
29        } else {
30            a[i] = a[m[i]] * a[i / m[i]];
31        }
32    }
33
34    let mut f = (1..=n)
35        .fold(ff.from_u64(1), |x, y| x * ff.from_u64(y as u64))
36        .inv();
37
38    for i in (0..=n).rev() {
39        a[i] *= f;
40        b[i] = f;
41        f *= ff.from_u64(i as u64);
42
43        if i % 2 == 1 {
44            b[i] = -b[i];
45        }
46    }
47
48    let mut ret = ntt.convolve(a, b);
49    ret.truncate(n + 1);
50    ret
51}
haar_lib/math/combinatorics/stirling_second.rs

haar_lib/math/combinatorics/
stirling_second.rs
