haar_lib/math/
stirling_first.rs

1//! 符号付き第一種スターリング数$s(n, 0), \dots, s(n, n)$を列挙する。
2//!
3//! $s(n,k)$ は $$x(x-1)\dots (x-(n-1)) = \sum_{k=0}^n s(n,k) x^k$$を満たす。
4use crate::math::ntt::NTT;
5use crate::math::polynomial::{Polynomial, PolynomialOperator};
6use crate::math::polynomial_taylor_shift::*;
7use crate::num::const_modint::*;
8
9/// 符号付き第一種スターリング数$s(n, 0), \dots, s(n, n)$を列挙する。
10///
11/// **Time complexity** $O(n \log n)$
12pub fn stirling_first<const P: u32, const PR: u32>(n: usize) -> Vec<ConstModInt<P>> {
13    let ff = ConstModIntBuilder;
14
15    let mut ret = Polynomial::<P>::from(vec![1]);
16    let ntt = NTT::<P, PR>::new();
17    let op = PolynomialOperator::new(&ntt);
18
19    let mut t: usize = 0;
20    let mut check = false;
21
22    for i in (0..usize::BITS).rev() {
23        if check {
24            let s = op.taylor_shift(ret.clone(), -ff.from_u64(t as u64));
25            ret = ntt.convolve(ret.into(), s.into()).into();
26            ret.as_mut().truncate(t * 2 + 1);
27            t *= 2;
28        }
29
30        if (n & (1 << i)) != 0 {
31            let a = vec![-ff.from_u64(t as u64), ff.from_u64(1)];
32            ret = ntt.convolve(ret.into(), a).into();
33            t += 1;
34
35            check = true;
36        }
37    }
38
39    ret.as_mut().truncate(n + 1);
40    ret.into()
41}
42
43#[cfg(test)]
44mod tests {
45    use crate::math::ntt::NTT998244353;
46
47    use super::*;
48
49    #[test]
50    fn test() {
51        const M: u32 = 998244353;
52        let ff = ConstModIntBuilder::<M>;
53        let ntt = NTT998244353::new();
54        let op = PolynomialOperator::new(&ntt);
55
56        let n = 100;
57        let mut ans = Polynomial::from(vec![ff.from_u64(1)]);
58
59        for i in 1..=n {
60            let res = stirling_first::<998244353, 3>(i);
61
62            ans = op.mul(
63                ans,
64                Polynomial::from(vec![-ff.from_u64(i as u64 - 1), 1.into()]),
65            );
66
67            assert_eq!(res, ans.as_ref());
68        }
69    }
70}