アルゴリズムと数学 042

https://atcoder.jp/contests/math-and-algorithm/tasks/abc172_d

よくある和の順番を変えるやつですね。約数の問題では特に多いです。

$\displaystyle S \equiv \sum_{K=1}^N{K\sum_{d|K}{1}}$

とおくと、

$\displaystyle S = \sum_{d=1}^N{\sum_{d|K}{K}}$

$K \equiv kd$ とおくと、

$\displaystyle S = \sum_{d=1}^N{\sum_{k=1}^{\lfloor \frac{N}{d} \rfloor}{kd}} = \sum_{d=1}^N{d\frac{\lfloor \frac{N}{d} \rfloor(\lfloor \frac{N}{d} \rfloor + 1)}{2}}$

となります。
これでも十分速いですが、dが大きいところでは $\lfloor \frac{N}{d} \rfloor$ が同じものが多いので、これをまとめると $O(N^{1/2})$ になります。

// Sum of Divisors
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}


//////////////////// process ////////////////////

fn f(N: i64) -> i64 {
    let r = (N as f64).sqrt() as i64;
    let s1 = (1..N/(r+1)+1).map(|d| d * (N/d) * (N/d + 1) / 2).sum::<i64>();
    let s2 = (1..r+1).map(|r| {
                        let first = N / (r + 1) + 1;
                        let last = N / r + 1;
                        (last - first) * (first + last - 1) * r * (r + 1) / 4 }).
                    sum::<i64>();
    s1 + s2
}

fn main() {
    let N: i64 = read();
    println!("{}", f(N))
}