haar_lib/math/convolution/
div_mul_transform.rs

1//! 約数・倍数 Zeta / Möbius 変換
2
3use std::ops::{Add, Sub};
4
5/// [`div_zeta`]の逆変換操作。
6pub fn div_mobius<T>(f: &mut [T])
7where
8    T: Copy + Sub<Output = T>,
9{
10    let n = f.len();
11    let mut check = vec![true; n];
12    for i in 2..n {
13        if check[i] {
14            for j in (1..=(n - 1) / i).rev() {
15                check[j * i] = false;
16                f[j * i] = f[j * i] - f[j];
17            }
18        }
19    }
20}
21
22/// $F_j = \sum_{j = 0 \pmod i} f_i$を満たす`F`を求める。
23pub fn div_zeta<T>(f: &mut [T])
24where
25    T: Copy + Add<Output = T>,
26{
27    let n = f.len();
28    let mut check = vec![true; n];
29    for i in 2..n {
30        if check[i] {
31            for j in (1..).take_while(|j| j * i < n) {
32                check[j * i] = false;
33                f[j * i] = f[j * i] + f[j];
34            }
35        }
36    }
37}
38
39/// [`mul_zeta`]の逆変換操作。
40pub fn mul_mobius<T>(f: &mut [T])
41where
42    T: Copy + Sub<Output = T>,
43{
44    let n = f.len();
45    let mut check = vec![true; n];
46    for i in 2..n {
47        if check[i] {
48            for j in (1..).take_while(|j| j * i < n) {
49                check[j * i] = false;
50                f[j] = f[j] - f[j * i];
51            }
52        }
53    }
54}
55
56/// $F_j = \sum_{i = 0 \pmod j} f_i$を満たす`F`を求める。
57pub fn mul_zeta<T>(f: &mut [T])
58where
59    T: Copy + Add<Output = T>,
60{
61    let n = f.len();
62    let mut check = vec![true; n];
63    for i in 2..n {
64        if check[i] {
65            for j in (1..=(n - 1) / i).rev() {
66                check[j * i] = false;
67                f[j] = f[j] + f[j * i];
68            }
69        }
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use crate::iter::collect::CollectVec;
76
77    use super::*;
78    use rand::Rng;
79
80    #[test]
81    fn test_div() {
82        let mut rng = rand::thread_rng();
83
84        let n = 500;
85        let m = 100000;
86
87        let f = (0..=n)
88            .map(|i| rng.gen_range(0..m) * i as u64 % m)
89            .collect_vec();
90
91        let mut g = vec![0; n + 1];
92        for j in 1..=n {
93            for i in 1..=n {
94                if j % i == 0 {
95                    g[j] += f[i];
96                }
97            }
98        }
99
100        let mut h = f.clone();
101        div_zeta(&mut h);
102        assert_eq!(h, g);
103
104        div_mobius(&mut h);
105        assert_eq!(h, f);
106    }
107
108    #[test]
109    fn test_mul() {
110        let mut rng = rand::thread_rng();
111
112        let n = 500;
113        let m = 100000;
114
115        let f = (0..=n)
116            .map(|i| rng.gen_range(0..m) * i as u64 % m)
117            .collect_vec();
118
119        let mut g = vec![0; n + 1];
120        for j in 1..=n {
121            for i in (j..=n).step_by(j) {
122                g[j] += f[i];
123            }
124        }
125
126        let mut h = f.clone();
127        mul_zeta(&mut h);
128        assert_eq!(h, g);
129
130        mul_mobius(&mut h);
131        assert_eq!(h, f);
132    }
133}
haar_lib/math/convolution/div_mul_transform.rs

haar_lib/math/convolution/
div_mul_transform.rs
