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