haar_lib/math/convolution/
conv_gcd.rs

1//! $\mathtt{a_{\gcd (i, j)}} = \sum \mathtt{f_{i}} * \mathtt{g_{j}}$を満たす`a`を求める。
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/gcd_convolution>
5use std::ops::{Add, Mul, Sub};
6
7/// $\mathtt{a_{\gcd (i, j)}} = \sum \mathtt{f_{i}} * \mathtt{g_{j}}$を満たす`a`を求める。
8pub fn convolution_gcd<T>(mut f: Vec<T>, mut g: Vec<T>) -> Vec<T>
9where
10    T: Copy + Add<Output = T> + Sub<Output = T> + Mul<Output = T>,
11{
12    assert_eq!(f.len(), g.len());
13
14    div_zeta(&mut f);
15    div_zeta(&mut g);
16
17    for (x, y) in f.iter_mut().zip(g.into_iter()) {
18        *x = *x * y;
19    }
20
21    div_mobius(&mut f);
22    f
23}
24
25fn div_mobius<T>(f: &mut [T])
26where
27    T: Copy + Sub<Output = T>,
28{
29    let n = f.len();
30    let mut check = vec![true; n];
31    for i in 2..n {
32        if check[i] {
33            for j in (1..).take_while(|j| j * i < n) {
34                check[j * i] = false;
35                f[j] = f[j] - f[j * i];
36            }
37        }
38    }
39}
40
41fn div_zeta<T>(f: &mut [T])
42where
43    T: Copy + Add<Output = T>,
44{
45    let n = f.len();
46    let mut check = vec![true; n];
47    for i in 2..n {
48        if check[i] {
49            for j in (1..=(n - 1) / i).rev() {
50                check[j * i] = false;
51                f[j] = f[j] + f[j * i];
52            }
53        }
54    }
55}
haar_lib/math/convolution/conv_gcd.rs

haar_lib/math/convolution/
conv_gcd.rs
