haar_lib/math/convolution/
subset_conv.rs

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

haar_lib/math/convolution/
subset_conv.rs
