前回のその１ではファウルハーバーの公式を２通りのやり方で求める方法を紹介しました。今回はその威力を発揮させるためにべき乗の和の剰余を求めます。
\begin{align}
S(p,n)=\sum_{k=1}^n k^p \pmod{m}
\end{align}
【例題 １】$S(5, 10^6) \pmod {(10^9+7)}$を求めよ
まずは素直に計算してみます
【コード#1】素直に計算
$n = 10^6$程度であれば１秒以内に求まります。もちろん剰余計算のべき乗はpow関数を使い、和もreduceを使って毎回剰余を取っています。
from functools import reduce
def S1(p, n, m):
def add_mod(a,b,m):return (a+b)%m
return reduce(lambda a, b : (a+b)%m, [pow(k, p, m) for k in range(1, n+1)])
p, n, m = 5, 10**6, 10**9+7
print(f"Answer : {S1(p, n, m)}")
# 419798333
【コード#2】ファウルハーバー係数を使う（逆行列）
ファウルハーバー係数を$\pmod {m}$で求める必要がありますが、最初のコードではsympyのbinomialとbernoulli関数は$\pmod {m}$をサポートしていないので、大幅に書き換える必要があります。
しかし行列を使ったコードでは逆行列を求めるときにinv_modを使うだけで$\pmod {m}$のファウルハーバー係数が求まります。
N = 6
m = 10**9+7
A, b = [], [1]
for j in range(1,N+1):
b = [1]+[(b[i]+b[i+1])%m for i in range(len(b)-1)]+[1] # binomia coefficient (Pascal's triangle)
A += [[b[i]*(-1)**(i+1+j) for i in range(len(b)-1)] + [0]*(N+1-len(b))] # change sign, if (i+1+j) is odd
A = sp.Matrix(A) # list -> Matrix
G = A.inv_mod(m) # inverse Matrix
for i in range(1, N):
print(i, list(G[i,0:(i+1)]))
# 1 [500000004, 500000004]
# 2 [166666668, 500000004, 333333336]
# 3 [0, 250000002, 500000004, 250000002]
# 4 [766666672, 0, 333333336, 500000004, 400000003]
# 5 [0, 916666673, 0, 416666670, 500000004, 166666668]
この５乗のファウルハーバー係数を使って例題１の答えを求めると、同じ答えが得られました。
コードは一見似ていますが、ループ回数$n=10^6$が$p+1=6$になっていることに注目して下さい。従って$n$がいくら大きくなっても計算量は変わらないことが分かります。
def S2(p, n, m):
return sum(list(G[p,0:(p+1)])[j]*pow(n, j+1, m) for j in range(p+1)) % m
p, n, m = 5, 10**6, 10**9+7
print(f"Answer: {S2(p, n, m)}")
# 186912818
ただし$p$が大きくなると$p \times p$の行列の逆行列を解く必要があるので、$p=100$ぐらいが限界です。それでは$p$も$n$も大きいときはどうすればよいかを次の例題で考えます。
【例題 ２】$S(10^3, 10^{14}) \pmod {10^9+7}$を求めよ
べき乗の和の剰余計算の公式を使う
いきなりですが、べき乗和の剰余計算では $m$が素数のとき以下の公式が成り立ちます。
\begin{equation}
\sum_{k=1}^{m}k^p \equiv
\begin{cases}
-1 \pmod{m}, & \text{if $m-1 \mid p$,} \\
0 \pmod{m}, & \text{if $m-1 \nmid p$,}
\end{cases}
\tag{2}
\end{equation}
本当かなと思うかもしれませんが、3から19までの素数
$m$に対して$p = [1, m]$でべき乗和を求めると$p=m-1$のときだけ$m-1 \equiv -1$となっているのが分かります。この証明は最後に乗せておきますので興味があれば読んで下さい。
import sympy as sp
for m in sp.primerange(3,20):
print(m, [sum(pow(k,p,m) for k in range(1,m+1))%m for p in range(1,m+1)])
# 3 [0, 2, 0]
# 5 [0, 0, 0, 4, 0]
# 7 [0, 0, 0, 0, 0, 6, 0]
# 11 [0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 0]
# 13 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 12, 0]
# 17 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 16, 0]
# 19 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 18, 0]
この公式を使うと式（１）は次のようになります。
\begin{align}
n'=n \bmod mとすると \\
\sum_{k=1}^n k^p \equiv \sum_{k=1}^{n'} k^p \pmod{m} \tag{3} \\
\end{align}
これで$n \lt m$にすることができます。
さらに$n' \gt m/2$ときは$m$までの残りを計算して$m$から引くことによって回答を得られます。
\begin{align}
\sum_{k=1}^n k^p \equiv m - \sum_{i=1}^{m-1-n'} (m-i)^p \equiv (-1)^{p+1}\sum_{i=1}^{m-1-n'} i^p \pmod{m} \tag{4}
\end{align}
【コード#3】べき乗の和の剰余計算の公式を使う
これをコードにすると、以下になり$p$と$n$の両方が大きい場合もなんとか対応できました。
def S3(p, n, m):
n %= m
if n > m-1-n:
return (-1)**(p+1)*S1(p, m-1-n, m)%m
return S1(p, n, m)
p, n, m = 10**3, 10**14, 10**9+7
print(f"Answer : {S3(p, n, m)}")
# Answer : 938341900
べき乗の和の剰余計算の公式の証明
原始根の定義と具体例（高校生向け）にあるように
- 任意の素数 p に対して原始根 r が存在する
- rがpの原始根のとき$\lbrace r, r^2, ... , r^{p-1} \rbrace \equiv \lbrace 1, 2, ... ,p-1 \rbrace \pmod p $
これを代入して、等比級数の公式を使うと、
\begin{align}
\sum_{i=1}^{p} i^k &\equiv \sum_{i=1}^{p-1} i^k \equiv \sum_{m=1}^{p-1}(r^m)^k \equiv \sum_{m=1}^{p-1}(r^{k})^{m} \\
&= \frac{1-(r^{k})^{p-1}}{1-r^{k}} \equiv 0 \pmod {p}
\end{align}
（開発環境：Google Colab）