haar_lib/math/
partition_number.rs

1//! 分割数$p(0), \dots, p(n)$を列挙する。
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/partition_function>
5use crate::{
6    math::{fps::inv::*, ntt::*, polynomial::*},
7    num::const_modint::*,
8};
9
10/// 分割数$p(0), \dots, p(n)$を列挙する。
11pub fn partition_number<const P: u32, const PR: u32>(n: usize) -> Vec<ConstModInt<P>> {
12    let ntt = NTT::<P, PR>::new();
13    let fps = PolynomialOperator::new(&ntt);
14
15    let ff = ConstModIntBuilder;
16    let mut f = vec![ff.from_u64(0); n + 1];
17    f[0] = ff.from_u64(1);
18
19    let m = ((1 + 24 * n).isqrt() - 1) / 6;
20    for i in 1..=m {
21        f[i * (3 * i + 1) / 2] += ff.from_i64(if i % 2 == 0 { 1 } else { -1 });
22    }
23
24    let m = ((1 + 24 * n).isqrt() + 1) / 6;
25    for i in 1..=m {
26        f[i * (3 * i - 1) / 2] += ff.from_i64(if i % 2 == 0 { 1 } else { -1 });
27    }
28
29    let f = Polynomial::from(f);
30    fps.fps_inv(f).unwrap().into()
31}
32
33#[cfg(test)]
34mod tests {
35    use super::*;
36
37    #[test]
38    fn test() {
39        let res = partition_number::<998244353, 3>(49);
40
41        let ans = [
42            1, 1, 2, 3, 5, 7, 11, 15, 22, 30, 42, 56, 77, 101, 135, 176, 231, 297, 385, 490, 627,
43            792, 1002, 1255, 1575, 1958, 2436, 3010, 3718, 4565, 5604, 6842, 8349, 10143, 12310,
44            14883, 17977, 21637, 26015, 31185, 37338, 44583, 53174, 63261, 75175, 89134, 105558,
45            124754, 147273, 173525,
46        ];
47
48        assert_eq!(res, ans.map(|x| x.into()));
49    }
50}
haar_lib/math/partition_number.rs

haar_lib/math/
partition_number.rs
