competitive/math/
lcm_convolve.rs

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

competitive/math/
lcm_convolve.rs
