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::polynomial::*;
7use crate::math::prime_mod::PrimeMod;
8use crate::num::const_modint::*;
9
10/// $\\#_p$ Subset sum
11pub fn number_of_subset_sum<P: PrimeMod>(s: Vec<usize>, t: usize) -> Vec<ConstModInt<P>> {
12    let ff = ConstModIntBuilder::<P>::new();
13
14    let mut c = vec![0; t + 1];
15    for x in s {
16        c[x] += 1;
17    }
18
19    let mut ret = vec![ff.from_u64(0); t + 1];
20
21    for (i, c) in c.into_iter().enumerate().skip(1) {
22        if c != 0 {
23            for j in (1..).take_while(|&j| i * j <= t) {
24                let k = j * i;
25                let x = ff.from_i64(if j % 2 == 1 { 1 } else { -1 })
26                    * ff.from_u64(i as u64)
27                    * ff.from_u64(k as u64).inv();
28                ret[k] += x * ff.from_u64(c as u64);
29            }
30        }
31    }
32
33    Polynomial::from(ret).fps_exp().unwrap().into()
34}