haar_lib/math/combinatorics/
partition_number.rs

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

haar_lib/math/combinatorics/
partition_number.rs
