ABC172-D Sum of Divisors 解説
\(O(\sqrt N )\) で解きます。
問題
リンクはこちら
概要
$$ \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[\frac{N}{i}] $$
を求めればよいですが、 \(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\) が平方数のケースなどでおかしくなることもあるので、注意が必要です。
提出コード
時間も余裕ですね。 \(O(N \log N)\) ってなんですか？
