haar_lib/algo/subset_sum/
count.rs

1//! 部分和問題 (和を達成する組み合わせ数を返す)
2use std::ops::Add;
3
4/// 部分和問題 (和を達成する組み合わせ数を返す)
5///
6/// **Time complexity** $O(nk)$
7///
8/// **Space complexity** $O(k)$
9pub fn subset_sum<T>(n: usize, k: usize, a: &[usize]) -> Vec<T>
10where
11    T: Copy + From<usize> + Add<Output = T>,
12{
13    assert!(a.len() == n);
14
15    let mut dp = vec![vec![T::from(0); k + 1]; 2];
16
17    dp[0][0] = T::from(1);
18
19    for (i, &x) in a.iter().enumerate() {
20        let cur = i & 1;
21        let next = (i + 1) & 1;
22        for j in 0..=k {
23            if j >= x {
24                dp[next][j] = dp[cur][j - x] + dp[cur][j];
25            } else {
26                dp[next][j] = dp[cur][j];
27            }
28        }
29    }
30
31    dp[n & 1].clone()
32}