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,
45    R::Additive: Invertible,
46{
47    type T = Vec<R::T>;
48    type F = Vec<R::T>;
49
50    fn length(t: &Self::T) -> usize {
51        t.len()
52    }
53
54    fn transform(mut t: Self::T, _len: usize) -> Self::F {
55        LcmConvolve::<R::Additive>::zeta_transform(&mut t);
56        t
57    }
58
59    fn inverse_transform(mut f: Self::F, _len: usize) -> Self::T {
60        LcmConvolve::<R::Additive>::mobius_transform(&mut f);
61        f
62    }
63
64    fn multiply(f: &mut Self::F, g: &Self::F) {
65        for (f, g) in f.iter_mut().zip(g) {
66            *f = R::mul(f, g);
67        }
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::{
75        algebra::{AddMulOperation, AdditiveOperation},
76        math::lcm,
77        rand,
78        tools::Xorshift,
79    };
80
81    const A: i64 = 100_000;
82
83    #[test]
84    fn test_lcm_convolve() {
85        let mut rng = Xorshift::new();
86
87        for m in 1..=300 {
88            rand!(rng, mut f: [-A..A; m]);
89            f[0] = 0;
90            let mut g = vec![0i64; m];
91            let h = f.clone();
92            for (s, f) in f.iter().enumerate().skip(1) {
93                for (t, g) in g.iter_mut().enumerate().skip(1) {
94                    if t % s == 0 {
95                        *g += f;
96                    }
97                }
98            }
99            LcmConvolve::<AdditiveOperation<i64>>::zeta_transform(&mut f);
100            assert_eq!(&f[1..], &g[1..]);
101            LcmConvolve::<AdditiveOperation<i64>>::mobius_transform(&mut f);
102            assert_eq!(&f[1..], &h[1..]);
103
104            rand!(rng, mut f: [-A..A; m], mut g: [-A..A; m]);
105            f[0] = 0;
106            g[0] = 0;
107            let mut h = vec![0i64; m];
108            for (i, f) in f.iter().enumerate().skip(1) {
109                for (j, g) in g.iter().enumerate().skip(1) {
110                    let k = lcm(i as _, j as _) as usize;
111                    if k < m {
112                        h[k] += f * g;
113                    }
114                }
115            }
116            let i = LcmConvolve::<AddMulOperation<i64>>::convolve(f, g);
117            assert_eq!(&h[1..], &i[1..]);
118        }
119    }
120}
competitive/math/lcm_convolve.rs

competitive/math/
lcm_convolve.rs
