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

competitive/math/
subset_convolve.rs
