haar_lib/math/convolution/
subset_conv.rs

1//! 添字部分集合畳み込み
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/subset_convolution>
5use crate::math::convolution::{mobius::*, zeta::*};
6use std::ops::{Add, Mul, Sub};
7
8/// $h_k = \sum_{i \lor j = k, i \land j = 0} f_i g_j$を満たす$h$を求める。
9///
10/// # Requirements
11/// `f.len()` = `g.len()`は2の累乗
12#[allow(clippy::needless_range_loop, clippy::manual_memcpy)]
13pub fn subset_convolution<T>(f: Vec<T>, g: Vec<T>) -> Vec<T>
14where
15    T: Copy + Default + Add<Output = T> + Sub<Output = T> + Mul<Output = T>,
16{
17    assert_eq!(f.len(), g.len());
18
19    let n = f.len();
20    assert!(n.is_power_of_two());
21
22    let k = n.trailing_zeros() as usize;
23
24    let mut f2 = vec![vec![T::default(); n]; k + 1];
25    let mut g2 = vec![vec![T::default(); n]; k + 1];
26
27    for j in 0..n {
28        let i = j.count_ones() as usize;
29        f2[i][j] = f[j];
30        g2[i][j] = g[j];
31    }
32
33    for i in 0..=k {
34        fast_zeta_subset(&mut f2[i]);
35        fast_zeta_subset(&mut g2[i]);
36    }
37
38    let mut h = vec![vec![T::default(); n]; k + 1];
39
40    for i in 0..=k {
41        for s in 0..=i {
42            for j in 0..n {
43                h[i][j] = h[i][j] + f2[s][j] * g2[i - s][j];
44            }
45        }
46    }
47
48    for i in 0..=k {
49        fast_mobius_subset(&mut h[i]);
50    }
51
52    (0..n)
53        .map(|j| {
54            let i = j.count_ones() as usize;
55            h[i][j]
56        })
57        .collect()
58}
haar_lib/math/convolution/subset_conv.rs

haar_lib/math/convolution/
subset_conv.rs
