haar_lib/math/
number_of_subset_sum.rs

1//! $\\#_p$ Subset sum
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/sharp_p_subset_sum>
5use crate::math::fps::exp::*;
6use crate::math::ntt::*;
7use crate::math::polynomial::*;
8use crate::num::const_modint::*;
9
10/// $\\#_p$ Subset sum
11pub fn number_of_subset_sum<const P: u32, const PR: u32>(
12    s: Vec<usize>,
13    t: usize,
14    ntt: &NTT<P, PR>,
15) -> Vec<ConstModInt<P>> {
16    let ff = ConstModIntBuilder;
17    let fps = PolynomialOperator::new(ntt);
18
19    let mut c = vec![0; t + 1];
20    for x in s {
21        c[x] += 1;
22    }
23
24    let mut ret = vec![ff.from_u64(0); t + 1];
25
26    for (i, c) in c.into_iter().enumerate().skip(1) {
27        if c != 0 {
28            for j in (1..).take_while(|&j| i * j <= t) {
29                let k = j * i;
30                let x = ff.from_i64(if j % 2 == 1 { 1 } else { -1 })
31                    * ff.from_u64(i as u64)
32                    * ff.from_u64(k as u64).inv();
33                ret[k] += x * ff.from_u64(c as u64);
34            }
35        }
36    }
37
38    fps.fps_exp(ret.into()).into()
39}
haar_lib/math/number_of_subset_sum.rs

haar_lib/math/
number_of_subset_sum.rs
