yukicoder No.2937 Sigma Plus Problem の O(log N) 解法

謝罪

元の解説文の一部に、作問者に失礼だと思われるような記述があるとのご指摘を受けましたので、該当する部分を全て削除いたしました。
誠に申し訳ございませんでした。

問題概要

問題文

ある正整数 $N$ が与えられます。
$\sum_{i=1}^Ni$ を求めてください。
ただし、 $N$ は $N\le 5\times10^9$ を満たす正整数とします。

解法

前提

まず、 $f(n)$ を以下のように定めます。$$f(n)=\sum_{i=1}^ni=1+2+\cdots+n$$この関数には、以下のような重要な性質があります。

• $f(n)$ の値がわかったら、 $f(n+1)$ の値を $O(1)$ で計算できる

$f(n+1)=f(n)+n+1$ で計算できます。

• $f(n)$ の値がわかったら、 $f(2n)$ の値も $O(1)$ で計算できる

これは、以下のように計算できます。

\begin{align}
f(2n)&=1+2+\cdots+n+(n+1)+(n+2)+\cdots+2n\\
&=2(1+2+\cdots+n)+n+n+\cdots+n\\
&=2f(n)+n^2
\end{align}

計算

$f(n)$ を求めるとき、以下のように再帰計算していくことを考えます。

• $n=1$ なら、答えは $1$ である。
• $n$ が奇数なら、 $f(n-1)$ を先に計算する。
• $n$ が偶数なら、 $f(n/2)$ を先に計算する。

この関数をそのまま実装すると、以下のようになります。

Python

def f(a):
    if a == 1:
        return 1
    elif a % 2:
        return f(a - 1) + a
    else:
        return 2 * f(a // 2) + (a // 2) * (a // 2)

計算量解析

そのままだと分かりづらいので、 $n$ を二進表記して考えます。

• $n$ が奇数なら、 $n$ の一番右の桁である $1$ を $0$ にする。
• $n$ が偶数なら、 $n$ の一番右の桁である $0$ を消す。

これを考慮すると、どのような $n$ の場合でも、高々 $2$ 回操作すれば桁数が $1$ つ減る、ということが言えます。
最終的に計算したいのは $f(N)$ であり、 $N$ を二進表記した時の桁数は $O(\log N)$ なので、全体の計算量も $O(\log N)$ です。

実装例

Python

def f(a):
    if a == 1:
        return 1
    elif a % 2:
        return f(a - 1) + a
    else:
        return 2 * f(a // 2) + (a // 2) * (a // 2)
N = int(input())
print(f(N))

C++

#include <bits/stdc++.h>
using namespace std;
using ull = unsigned long long;
ull f(ull a) {
    if (a == 1) {
        return 1;
    } else if (a % 2) {
        return f(a - 1) + a;
    } else {
        return 2 * f(a / 2) + (a / 2) * (a / 2);
    }
}
int main() {
    ull N;
    cin >> N;
    cout << f(N) << '\n';
}

おわり。