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>(
13    n: usize,
14    ntt: &NTT<P, PR>,
15) -> Vec<ConstModInt<P>> {
16    let ff = ConstModIntBuilder;
17
18    let mut ret = Polynomial::<P>::from(vec![1]);
19    let op = PolynomialOperator::new(ntt);
20
21    let mut t: usize = 0;
22    let mut check = false;
23
24    for i in (0..usize::BITS).rev() {
25        if check {
26            let s = op.taylor_shift(ret.clone(), -ff.from_u64(t as u64));
27            ret = ntt.convolve(ret.into(), s.into()).into();
28            ret.as_mut().truncate(t * 2 + 1);
29            t *= 2;
30        }
31
32        if (n & (1 << i)) != 0 {
33            let a = vec![-ff.from_u64(t as u64), ff.from_u64(1)];
34            ret = ntt.convolve(ret.into(), a).into();
35            t += 1;
36
37            check = true;
38        }
39    }
40
41    ret.as_mut().truncate(n + 1);
42    ret.into()
43}
44
45#[cfg(test)]
46mod tests {
47    use crate::math::ntt::NTT998244353;
48
49    use super::*;
50
51    #[test]
52    fn test() {
53        const M: u32 = 998244353;
54        let ff = ConstModIntBuilder::<M>;
55        let ntt = NTT998244353::new();
56        let op = PolynomialOperator::new(&ntt);
57
58        let n = 100;
59        let mut ans = Polynomial::from(vec![ff.from_u64(1)]);
60
61        for i in 1..=n {
62            let res = stirling_first(i, &ntt);
63
64            ans = op.mul(
65                ans,
66                Polynomial::from(vec![-ff.from_u64(i as u64 - 1), 1.into()]),
67            );
68
69            assert_eq!(res, ans.as_ref());
70        }
71    }
72}