haar_lib/math/
stirling_second_fixed_k.rs

1//! 第二種スターリング数$S(0, k), \dots, S(n, k)$を列挙する。
2use crate::math::factorial::FactorialTable;
3use crate::math::fps::pow::FpsPow;
4use crate::math::ntt::NTT;
5use crate::math::polynomial::PolynomialOperator;
6use crate::num::const_modint::*;
7
8/// 第二種スターリング数$S(0, k), \dots, S(n, k)$を列挙する。
9pub fn stirling_second_fixed_k<const P: u32, const PR: u32>(
10    n: usize,
11    k: usize,
12) -> Vec<ConstModInt<P>> {
13    assert!(k <= n);
14
15    let ntt = NTT::<P, PR>::new();
16    let fps = PolynomialOperator::new(&ntt);
17    let ft = FactorialTable::new(n, ConstModIntBuilder);
18
19    let mut ret = vec![ConstModInt::new(0); n + 1];
20
21    for i in 1..=n {
22        ret[i] = ft.inv_facto(i);
23    }
24
25    ret = fps.fps_pow(ret.into(), k as u64).unwrap().into();
26
27    for i in k..=n {
28        ret[i] *= ft.inv_facto(k) * ft.facto(i);
29    }
30
31    ret
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37
38    use crate::math::stirling_second_table::stirling_second_table;
39
40    #[test]
41    fn test() {
42        let n = 100;
43        let ans = stirling_second_table(n, ConstModIntBuilder::<998244353>);
44
45        for k in 0..=n {
46            assert_eq!(
47                stirling_second_fixed_k::<998244353, 3>(n, k),
48                ans.iter().map(|a| a[k]).collect::<Vec<_>>()
49            );
50        }
51    }
52}