competitive/math/
subset_convolve.rs

1use super::{BitwiseorConvolve, ConvolveSteps, Invertible, Ring};
2use std::marker::PhantomData;
3
4pub struct SubsetConvolve<M> {
5    _marker: PhantomData<fn() -> M>,
6}
7
8impl<R> ConvolveSteps for SubsetConvolve<R>
9where
10    R: Ring,
11    R::T: PartialEq,
12    R::Additive: Invertible,
13{
14    type T = Vec<R::T>;
15    type F = Vec<Vec<R::T>>;
16
17    fn length(t: &Self::T) -> usize {
18        t.len()
19    }
20
21    fn transform(t: Self::T, len: usize) -> Self::F {
22        let k = len.trailing_zeros() as usize;
23        let mut f = vec![vec![R::zero(); len]; k + 1];
24        for (i, t) in t.iter().enumerate() {
25            let f = &mut f[i.count_ones() as usize][i];
26            *f = R::add(f, t);
27        }
28        for f in f.iter_mut() {
29            BitwiseorConvolve::<R::Additive>::zeta_transform(f);
30        }
31        f
32    }
33
34    fn inverse_transform(mut f: Self::F, len: usize) -> Self::T {
35        for f in f.iter_mut() {
36            BitwiseorConvolve::<R::Additive>::mobius_transform(f);
37        }
38        let mut t = vec![R::zero(); len];
39        for (i, t) in t.iter_mut().enumerate() {
40            *t = R::add(t, &f[i.count_ones() as usize][i]);
41        }
42        t
43    }
44
45    fn multiply(f: &mut Self::F, g: &Self::F) {
46        for i in (0..f.len()).rev() {
47            let (lf, rf) = f.split_at_mut(i);
48            for (x, y) in rf[0].iter_mut().zip(&g[0]) {
49                *x = R::mul(x, y);
50            }
51            for (x, y) in lf.iter().rev().zip(g.iter().skip(1)) {
52                for ((x, y), z) in x.iter().zip(y).zip(&mut rf[0]) {
53                    *z = R::add(z, &R::mul(x, y));
54                }
55            }
56        }
57    }
58
59    fn convolve(a: Self::T, b: Self::T) -> Self::T {
60        assert_eq!(a.len(), b.len());
61        let len = a.len();
62        let same = a == b;
63        let mut a = Self::transform(a, len);
64        let b = if same {
65            a.clone()
66        } else {
67            Self::transform(b, len)
68        };
69        Self::multiply(&mut a, &b);
70        Self::inverse_transform(a, len)
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use crate::{algebra::AddMulOperation, rand, tools::Xorshift};
78
79    const A: i64 = 100_000;
80
81    #[test]
82    fn test_subset_convolve() {
83        let mut rng = Xorshift::new();
84
85        for k in 0..12 {
86            let n = 1 << k;
87            rand!(rng, f: [-A..A; n], g: [-A..A; n]);
88            let mut h = vec![0i64; n];
89            for i in 0..n {
90                for j in 0..n {
91                    if i & j == 0 {
92                        h[i | j] += f[i] * g[j];
93                    }
94                }
95            }
96            let i = SubsetConvolve::<AddMulOperation<i64>>::convolve(f, g);
97            assert_eq!(h, i);
98        }
99    }
100}
competitive/math/subset_convolve.rs

competitive/math/
subset_convolve.rs
