haar_lib/math/
stirling_first_fixed_k.rs1use 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
9pub 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}