haar_lib/math/
totient_sum.rs

1//! トーシェント関数の総和
2//!
3//! # References
4//! - <https://yukicoder.me/wiki/sum_totient>
5//!
6//! # Problems
7//! - <https://judge.yosupo.jp/problem/sum_of_totient_function>
8
9use crate::math::totient::totient_table;
10use crate::num::ff::*;
11use std::collections::HashMap;
12
13/// トーシェント関数の総和
14pub fn totient_sum<Modulo: FF>(n: u64, m: Modulo) -> Modulo::Element
15where
16    Modulo::Element: FFElem + Copy,
17{
18    let k = (n as f64).powf(0.66) as u64;
19
20    let mut memo1 = vec![m.from_u64(0); k as usize + 1];
21    let table = totient_table(k as usize);
22    let mut sum = m.from_u64(0);
23    for i in 1..=k as usize {
24        sum += m.from_u64(table[i]);
25        memo1[i] = sum;
26    }
27
28    let mut memo2 = HashMap::new();
29
30    rec(n, &m, &memo1, &mut memo2)
31}
32
33fn rec<Modulo: FF>(
34    x: u64,
35    m: &Modulo,
36    memo1: &[Modulo::Element],
37    memo2: &mut HashMap<u64, Modulo::Element>,
38) -> Modulo::Element
39where
40    Modulo::Element: FFElem + Copy,
41{
42    if let Some(y) = memo1.get(x as usize) {
43        return *y;
44    }
45    if let Some(y) = memo2.get(&x) {
46        return *y;
47    }
48
49    let mut ret = if x % 2 == 0 {
50        m.from_u64(x / 2) * m.from_u64(x + 1)
51    } else {
52        m.from_u64(x) * m.from_u64(x.div_ceil(2))
53    };
54
55    let s = (x as f64).sqrt() as u64;
56
57    for i in 2..=s {
58        if x / i > s {
59            ret -= rec(x / i, m, memo1, memo2);
60        }
61    }
62
63    for i in 1..=s {
64        let t = m.from_u64(x / i - x / (i + 1));
65        ret -= rec(i, m, memo1, memo2) * t;
66    }
67
68    memo2.insert(x, ret);
69    ret
70}
haar_lib/math/totient_sum.rs

haar_lib/math/
totient_sum.rs
