ABC172-D Sum of Divisors 解説
$O(\sqrt N)$ で解きます。
問題
AtCoder Beginner Contest 172 | D - Sum of Divisors
概要
$$ \sum _ {i=1} ^ {N} i \cdot \mathit{d}(i) $$
を求めなさい。
解法
公式editorialの $O(N)$ 解法を高速化することを考えたいです。結局、 $S[k] := \sum _ {i=1} ^ {k} i$ として
$$ \sum _ {i=1} ^ {N} i \cdot S \lbrack \frac{N}{i} \rbrack $$
を求めればよいですが、 $i$ が後半になるにつれ、 $\frac{N}{i}$ の値は同じものが連続するようになってきます。
そこで、 $S[\frac{N}{i}]$ の値が等しくなるような $i$ についてまとめて計算することを考えます。その準備として、 $M := \sqrt N$ とおき、
- 前半 ( $i < M$ )
- 後半 ( $i \geq M$ )
の範囲に分けて考えます。
前半 ( $i < M$ )
公式解説通り一つずつ足し合わせます。計算量は $O(M)$ です。
後半 ( $i \geq M$ )
$M$の定義から、この範囲では $\frac{N}{i} \leq M$ が成り立ちます。そこで、 $\frac{N}{i}$ の側を全探索してみましょう。
$\frac{N}{i} = k$ となるようなiの範囲を計算すると $\frac{N}{k+1} + 1 \leq i \leq \frac{N}{k}$ であるので、
$$ \sum _ {i = \frac{N}{k+1} + 1} ^ {\frac{N}{k}} i \cdot S[k] $$
が高速で求められると嬉しいです。変形すると、
$$ \begin{aligned} \sum _ {i = \frac{N}{k+1} + 1} ^ {\frac{N}{k}} i \cdot S[k] &= S[k] \cdot\sum _ {i = \frac{N}{k+1} + 1} ^ {\frac{N}{k}} i \\ &= S[k] \cdot (S[\frac{N}{k}] - S[\frac{N}{k+1}]) \end{aligned} $$
であり、これは $O(1)$ で求められます。
あとは、これを $k = 1, 2, ..., M$ について足し上げればよいです。計算量は $O(M)$ です。
$S$ の前計算について
累積和の要領で $S$ を前計算しようとすると $O(N)$ かかってしまい本末転倒です。これは簡単な等差数列の和なので、配列で持つのをやめて毎回
$$ S(k) = \frac{k(k+1)}{2} $$
と計算してしまえば前計算が必要ありません。
以上より、全体の答えを $O(M) = O(\sqrt N)$ で求めることができました。
実装は前半後半の境界（$M$）付近がごちゃごちゃになりがちなので、前半部分を求める時は $M$ まで求めてしまうのではなく、後半部分がどこまでカバーできるのかを計算して、それに合わせて計算範囲を決めると楽です。また、 $M$ の計算方法によっては $N$ が極端に大きい・小さいケースや $N$ が平方数のケースなどでおかしくなることもあるので、注意が必要です。
実装
def S(k):
return k * (k + 1) // 2
N = int(input())
M = int((N+1)**.5)
ans = 0
# 前半
for i in range(1, N // (M+1) + 1):
ans += i * S(N // i)
# 後半
for k in range(1, M+1):
ans += S(k) * (S(N//k) - S(N//(k+1)))
print(ans)
時間も余裕ですね。 $O(N \log N)$ ってなんですか？