プログラミング問題『累積和（Σx^k）』の解き方

Aru

この記事では、累積和の公式である$\sigma X^k = X^0 + X^1 + X^2 + … + X^n$を求める方法を解説します。この問題は、数学的アプローチを用いることで容易に解けるものですが、競技プログラミングのコンテストで出題されると悩む問題でもあります。

その他のAtCoderに役立つ記事の一覧

AtCoderで役立つ記事一覧（Python, Go言語）

Contents

はじめに
- 解きたかった問題
- 一般化
数学的な解法
Pythonで問題を解いてみる
まとめ

はじめに

今回、この内容を記事にしたのは、ある問題を解く時に悩んだからです。

解きたかった問題

解きたかった問題は、アルゴリズム×数学が基礎からしっかり身につく本の、053の問題「053 – Sum of 4^N」です。

問題は以下のようなものです。

問題

$4^0 + 4^1 + …+4^N$を求めなさい

Nは$N \leq 10^{18}$と非常に大きいため、$4^0, 4^1,…,4^N$を順番に計算していくのでは間に合いません。

結論ですが、この問題は、以下の式の結果が答えになります。

解答

$$\frac{4^{n+1} – 1}{3}$$

この問題には、公式の解説がなかったので自力で解けなければ、どう考えて良いかわかりません。ということで、自分がどのように解いたのかを記事にしてみました。

一般化

上記の問題ですが、とりあえず一般化した以下の形式で考えたいと思います。

$$
x^0 + x^1 + … x^n = \sum_{k=0}^{n}{x^k}
$$

数学的な解法

まず、上の式を$s$とおきます。

$$\begin{eqnarray}
s = \sum_{k=0}^{n}{x^k} &=& x^0 + x^1 + x^2 + … + x^n\\
&=& 1 + x^1 + x^2 + … + x^n
\end{eqnarray}$$

ここで、$x\times s$を考えると、以下のようになります。

$$\begin{eqnarray}
xs &= & x(1 + x^1 + x^2 + … + x^n)\\
&=& x^1 + x^2 + x^3 + … + x^{n+1}\\
&=& (1 + x^1 + x^2 + x^3 + … + x^n) -1 + x^{n+1}\\
&=& s – 1 + x^{n+1}
\end{eqnarray}$$

これを$s$について解きます。

$$\begin{eqnarray}
(x-1)s = x^{n+1} – 1\\
s = \frac{x^{n+1}-1}{x-1}
\end{eqnarray}$$

以上でsを求めることができました。あとは、これを計算するだけです。

結果

$$\sum_{k=0}^{n}{x^k} = \frac{x^{n+1}-1}{x-1}$$

Pythonで問題を解いてみる

pythonの場合、powがmod付きで計算できるので楽です。

まず、$4^{n+1}$を計算し、次に-1しています。

-1した結果が負になった場合は（mod演算なので0になる可能性がある）$10^9+7$だけ足しています。

div3は、逆元の計算です（フェルマーの小定理）。合同式で割り算の計算をする場合は、少し工夫が必要になります。1/3（逆元）を求めたので、これを乗じたら結果となります。

n = int(input())

mod = 10**9 + 7
ans = pow(4, n+1, mod)
ans -= 1
if ans < 0 : ans += mod
div3 = pow(3,  mod-2, mod)
ans = ans * div3 % mod
print(ans)