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>(
12    n: usize,
13    ntt: &NTT<P, PR>,
14) -> Vec<ConstModInt<P>> {
15    let fps = PolynomialOperator::new(ntt);
16
17    let ff = ConstModIntBuilder;
18    let mut f = vec![ff.from_u64(0); n + 1];
19    f[0] = ff.from_u64(1);
20
21    let m = ((1 + 24 * n).isqrt() - 1) / 6;
22    for i in 1..=m {
23        f[i * (3 * i + 1) / 2] += ff.from_i64(if i % 2 == 0 { 1 } else { -1 });
24    }
25
26    let m = ((1 + 24 * n).isqrt() + 1) / 6;
27    for i in 1..=m {
28        f[i * (3 * i - 1) / 2] += ff.from_i64(if i % 2 == 0 { 1 } else { -1 });
29    }
30
31    let f = Polynomial::from(f);
32    fps.fps_inv(f).into()
33}
34
35#[cfg(test)]
36mod tests {
37    use crate::math::ntt::NTT998244353;
38
39    use super::*;
40
41    #[test]
42    fn test() {
43        let ntt = NTT998244353::new();
44
45        let res = partition_number(49, &ntt);
46
47        let ans = [
48            1, 1, 2, 3, 5, 7, 11, 15, 22, 30, 42, 56, 77, 101, 135, 176, 231, 297, 385, 490, 627,
49            792, 1002, 1255, 1575, 1958, 2436, 3010, 3718, 4565, 5604, 6842, 8349, 10143, 12310,
50            14883, 17977, 21637, 26015, 31185, 37338, 44583, 53174, 63261, 75175, 89134, 105558,
51            124754, 147273, 173525,
52        ];
53
54        assert_eq!(res, ans.map(|x| x.into()));
55    }
56}
haar_lib/math/partition_number.rs

haar_lib/math/
partition_number.rs
