haar_lib/math/
multiplicative.rs

1//! 乗法的関数
2
3use std::ops::Mul;
4
5/// 乗法的関数を列挙する。
6///
7/// `f(p, k)`は乗法的関数の$p^k$での値を返す。
8pub fn multiplicative_table<T, F>(n: usize, one: T, f: F) -> Vec<T>
9where
10    T: Copy + Mul<Output = T>,
11    F: Fn(u64, u32) -> T,
12{
13    let mut ret = vec![one; n + 1];
14    let mut p = vec![(0, 0, 0); n + 1];
15
16    for i in 2..=n {
17        if p[i] == (0, 0, 0) {
18            for j in (i..=n).step_by(i) {
19                if p[j] == (0, 0, 0) {
20                    let mut k = 0;
21                    let mut m = j;
22
23                    while m % i == 0 {
24                        m /= i;
25                        k += 1;
26                    }
27
28                    p[j] = (m, i, k);
29                }
30            }
31        }
32    }
33
34    for i in 2..=n {
35        let (m, j, k) = p[i];
36        ret[i] = ret[m] * f(j as u64, k);
37    }
38
39    ret
40}
41
42#[cfg(test)]
43mod tests {
44    use super::*;
45    use crate::{math::totient::totient_table, timer};
46
47    #[test]
48    fn totient() {
49        let f = |p: u64, k: u32| p.pow(k) - p.pow(k - 1);
50
51        for n in [1, 10, 100, 1000, 10000, 100000, 1000000] {
52            let res = timer! {n, {
53                multiplicative_table(n, 1, f)
54            }};
55
56            let ans = totient_table(n);
57
58            assert_eq!(res[1..], ans[1..]);
59        }
60    }
61
62    #[test]
63    fn mobius() {
64        let f = |_p: u64, k: u32| if k == 1 { -1 } else { 0 };
65
66        for n in [1, 10, 100, 1000, 10000, 100000, 1000000] {
67            let _res = timer! {n, {
68                multiplicative_table(n, 1, f)
69            }};
70        }
71    }
72}
haar_lib/math/multiplicative.rs

haar_lib/math/
multiplicative.rs
