haar_lib/math/
multiplicative.rs

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

haar_lib/math/
multiplicative.rs
