competitive/math/
bitwiseand_convolve.rs

1use super::{ConvolveSteps, Group, Invertible, Monoid, Ring, bitwise_transform};
2use std::marker::PhantomData;
3
4pub struct BitwiseandConvolve<M> {
5    _marker: PhantomData<fn() -> M>,
6}
7
8impl<M> BitwiseandConvolve<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, |x, y| *x = M::operate(x, y));
15    }
16}
17
18impl<G> BitwiseandConvolve<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, |x, y| *x = G::rinv_operate(x, y));
25    }
26}
27
28impl<R> ConvolveSteps for BitwiseandConvolve<R>
29where
30    R: Ring<T: PartialEq, Additive: Invertible>,
31{
32    type T = Vec<R::T>;
33    type F = Vec<R::T>;
34
35    fn length(t: &Self::T) -> usize {
36        t.len()
37    }
38
39    fn transform(mut t: Self::T, _len: usize) -> Self::F {
40        BitwiseandConvolve::<R::Additive>::zeta_transform(&mut t);
41        t
42    }
43
44    fn inverse_transform(mut f: Self::F, _len: usize) -> Self::T {
45        BitwiseandConvolve::<R::Additive>::mobius_transform(&mut f);
46        f
47    }
48
49    fn multiply(f: &mut Self::F, g: &Self::F) {
50        for (f, g) in f.iter_mut().zip(g) {
51            *f = R::mul(f, g);
52        }
53    }
54
55    fn convolve(a: Self::T, b: Self::T) -> Self::T {
56        assert_eq!(a.len(), b.len());
57        let len = a.len();
58        let same = a == b;
59        let mut a = Self::transform(a, len);
60        if same {
61            for a in a.iter_mut() {
62                *a = R::mul(a, a);
63            }
64        } else {
65            let b = Self::transform(b, len);
66            Self::multiply(&mut a, &b);
67        }
68        Self::inverse_transform(a, len)
69    }
70}
71
72pub struct OnlineSupersetZetaTransform<M>
73where
74    M: Monoid,
75{
76    data: Vec<M::T>,
77}
78
79impl<M> OnlineSupersetZetaTransform<M>
80where
81    M: Monoid,
82{
83    pub fn new(size: usize) -> Self {
84        Self {
85            data: Vec::with_capacity(size),
86        }
87    }
88
89    pub fn push(&mut self, x: M::T) {
90        let i = self.data.len();
91        self.data.push(x);
92        let mut k = 0;
93        while i >> k & 1 == 1 {
94            let size = 1 << k;
95            let chunk = &mut self.data[i - (size * 2 - 1)..];
96            let (x, y) = chunk.split_at_mut(size);
97            for (x, y) in x.iter_mut().zip(y.iter()) {
98                M::operate_assign(x, y);
99            }
100            k += 1;
101        }
102    }
103
104    pub fn get(&self, index: usize) -> M::T {
105        let mut cur = self.data.len();
106        let mut s = M::unit();
107        while cur != 0 {
108            let d = cur.trailing_zeros() as usize;
109            let base = cur ^ (1 << d);
110            if (!base & index) >> d == 0 {
111                let mask = (1 << d) - 1;
112                s = M::operate(&s, &self.data[base ^ ((index ^ base) & mask)]);
113            }
114            cur = base;
115        }
116        s
117    }
118}
119
120pub struct OnlineSupersetMobiusTransform<M>
121where
122    M: Group,
123{
124    data: Vec<M::T>,
125}
126
127impl<M> OnlineSupersetMobiusTransform<M>
128where
129    M: Group,
130{
131    pub fn new(size: usize) -> Self {
132        Self {
133            data: Vec::with_capacity(size),
134        }
135    }
136
137    pub fn push(&mut self, x: M::T) {
138        let i = self.data.len();
139        self.data.push(x);
140        let mut k = 0;
141        while i >> k & 1 == 1 {
142            let size = 1 << k;
143            let chunk = &mut self.data[i - (size * 2 - 1)..];
144            let (x, y) = chunk.split_at_mut(size);
145            for (x, y) in x.iter_mut().zip(y.iter()) {
146                M::rinv_operate_assign(x, y);
147            }
148            k += 1;
149        }
150    }
151
152    pub fn get(&self, index: usize) -> M::T {
153        let mut cur = self.data.len();
154        let mut s = M::unit();
155        while cur != 0 {
156            let d = cur.trailing_zeros() as usize;
157            let base = cur ^ (1 << d);
158            if (!base & index) >> d == 0 {
159                let mask = (1 << d) - 1;
160                let j = base ^ ((index ^ base) & mask);
161                if ((index ^ base) >> d).count_ones() & 1 == 1 {
162                    s = M::rinv_operate(&s, &self.data[j]);
163                } else {
164                    s = M::operate(&s, &self.data[j]);
165                }
166            }
167            cur = base;
168        }
169        s
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use crate::{
177        algebra::{AddMulOperation, AdditiveOperation},
178        rand,
179        tools::Xorshift,
180    };
181
182    const A: i64 = 100_000;
183
184    #[test]
185    fn test_bitwiseand_convolve() {
186        let mut rng = Xorshift::default();
187
188        for k in 0..12 {
189            let n = 1 << k;
190            rand!(rng, mut f: [-A..A; n]);
191            let mut g = vec![0i64; n];
192            let h = f.clone();
193            for (s, f) in f.iter().enumerate() {
194                for (t, g) in g.iter_mut().enumerate() {
195                    if s | t == s {
196                        *g += f;
197                    }
198                }
199            }
200            BitwiseandConvolve::<AdditiveOperation<i64>>::zeta_transform(&mut f);
201            assert_eq!(f, g);
202            BitwiseandConvolve::<AdditiveOperation<i64>>::mobius_transform(&mut f);
203            assert_eq!(f, h);
204
205            rand!(rng, f: [-A..A; n], g: [-A..A; n]);
206            let mut h = vec![0i64; n];
207            for i in 0..n {
208                for j in 0..n {
209                    h[i & j] += f[i] * g[j];
210                }
211            }
212            let i = BitwiseandConvolve::<AddMulOperation<i64>>::convolve(f, g);
213            assert_eq!(h, i);
214        }
215    }
216
217    #[test]
218    fn test_online_superset_zeta_transform() {
219        let mut rng = Xorshift::default();
220        for k in 0..12 {
221            let n = 1 << k;
222            rand!(rng, f: [-A..A; n]);
223            let mut ost = OnlineSupersetZetaTransform::<AdditiveOperation<i64>>::new(0);
224            for i in 0..n {
225                ost.push(f[i]);
226                let mut g: Vec<_> = f[..=i]
227                    .iter()
228                    .cloned()
229                    .chain(std::iter::repeat_n(0, n - i - 1))
230                    .collect();
231                BitwiseandConvolve::<AdditiveOperation<i64>>::zeta_transform(&mut g);
232                for (j, &g) in g.iter().enumerate() {
233                    assert_eq!(ost.get(j), g);
234                }
235            }
236        }
237    }
238
239    #[test]
240    fn test_online_superset_mobius_transform() {
241        let mut rng = Xorshift::default();
242        for k in 0..12 {
243            let n = 1 << k;
244            rand!(rng, f: [-A..A; n]);
245            let mut ost = OnlineSupersetMobiusTransform::<AdditiveOperation<i64>>::new(0);
246            for i in 0..n {
247                ost.push(f[i]);
248                let mut g: Vec<_> = f[..=i]
249                    .iter()
250                    .cloned()
251                    .chain(std::iter::repeat_n(0, n - i - 1))
252                    .collect();
253                BitwiseandConvolve::<AdditiveOperation<i64>>::mobius_transform(&mut g);
254                for (j, &g) in g.iter().enumerate() {
255                    assert_eq!(ost.get(j), g);
256                }
257            }
258        }
259    }
260}
competitive/math/bitwiseand_convolve.rs

competitive/math/
bitwiseand_convolve.rs
