competitive/math/
bitwiseor_convolve.rs

1use super::{ConvolveSteps, Group, Invertible, Monoid, Ring, bitwise_transform};
2use std::marker::PhantomData;
3
4pub struct BitwiseorConvolve<M> {
5    _marker: PhantomData<fn() -> M>,
6}
7
8impl<M> BitwiseorConvolve<M>
9where
10    M: Monoid,
11{
12    /// $$g(m) = \sum_{n \mid m}f(n)$$
13    pub fn zeta_transform(f: &mut [M::T]) {
14        bitwise_transform(f, |y, x| *x = M::operate(x, y));
15    }
16}
17
18impl<G> BitwiseorConvolve<G>
19where
20    G: Group,
21{
22    /// $$f(m) = \sum_{n \mid m}h(n)$$
23    pub fn mobius_transform(f: &mut [G::T]) {
24        bitwise_transform(f, |y, x| *x = G::rinv_operate(x, y));
25    }
26}
27
28impl<R> ConvolveSteps for BitwiseorConvolve<R>
29where
30    R: Ring,
31    R::T: PartialEq,
32    R::Additive: Invertible,
33{
34    type T = Vec<R::T>;
35    type F = Vec<R::T>;
36
37    fn length(t: &Self::T) -> usize {
38        t.len()
39    }
40
41    fn transform(mut t: Self::T, _len: usize) -> Self::F {
42        BitwiseorConvolve::<R::Additive>::zeta_transform(&mut t);
43        t
44    }
45
46    fn inverse_transform(mut f: Self::F, _len: usize) -> Self::T {
47        BitwiseorConvolve::<R::Additive>::mobius_transform(&mut f);
48        f
49    }
50
51    fn multiply(f: &mut Self::F, g: &Self::F) {
52        for (f, g) in f.iter_mut().zip(g) {
53            *f = R::mul(f, g);
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        if same {
63            for a in a.iter_mut() {
64                *a = R::mul(a, a);
65            }
66        } else {
67            let b = Self::transform(b, len);
68            Self::multiply(&mut a, &b);
69        }
70        Self::inverse_transform(a, len)
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use crate::{
78        algebra::{AddMulOperation, AdditiveOperation},
79        rand,
80        tools::Xorshift,
81    };
82
83    const A: i64 = 100_000;
84
85    #[test]
86    fn test_bitwiseor_convolve() {
87        let mut rng = Xorshift::new();
88
89        for k in 0..12 {
90            let n = 1 << k;
91            rand!(rng, mut f: [-A..A; n]);
92            let mut g = vec![0i64; n];
93            let h = f.clone();
94            for (s, f) in f.iter().enumerate() {
95                for (t, g) in g.iter_mut().enumerate() {
96                    if s | t == t {
97                        *g += f;
98                    }
99                }
100            }
101            BitwiseorConvolve::<AdditiveOperation<i64>>::zeta_transform(&mut f);
102            assert_eq!(f, g);
103            BitwiseorConvolve::<AdditiveOperation<i64>>::mobius_transform(&mut f);
104            assert_eq!(f, h);
105
106            rand!(rng, f: [-A..A; n], g: [-A..A; n]);
107            let mut h = vec![0i64; n];
108            for i in 0..n {
109                for j in 0..n {
110                    h[i | j] += f[i] * g[j];
111                }
112            }
113            let i = BitwiseorConvolve::<AddMulOperation<i64>>::convolve(f, g);
114            assert_eq!(h, i);
115        }
116    }
117}
competitive/math/bitwiseor_convolve.rs

competitive/math/
bitwiseor_convolve.rs
