dsalgo/
fast_zeta_transform_for_multiples_with_instant_func.rs

1use std::ops::*;
2
3use crate::sieve_of_eratosthenes_enumerate_primes_usize::enumerate_primes;
4
5pub fn fast_zeta_multiples<T, F>(
6    op: F,
7    mut f: Vec<T>,
8) -> Vec<T>
9where
10    T: Clone,
11    F: Fn(T, T) -> T,
12{
13    let n = f.len();
14
15    for p in enumerate_primes(n) {
16        for i in (1..(n - 1) / p + 1).rev() {
17            f[i] = op(f[i].clone(), f[i * p].clone());
18        }
19    }
20
21    f
22}
23
24#[cfg(test)]
25
26mod tests {
27
28    use super::*;
29
30    #[test]
31
32    fn test() {
33        use crate::number_of_multiples_table_with_multiple_zeta_usize::*;
34
35        let n = 1 << 15;
36
37        let mut a = vec![1; n];
38
39        a[0] = 0;
40
41        let f = |a, b| a + b;
42
43        let res = fast_zeta_multiples(f, a);
44
45        assert_eq!(res, number_of_multiples(n));
46    }
47}