高速な冪乗和 ( Σ 計算 ) + mod の計算について

まえがき

この記事は投稿者(NokonoKotlin)の個人サイトの記事から Qiita 用に移植 & 加筆したものです。 (この文言は記事が剽窃でないことを周知するためのものであり、個人サイトの宣伝ではないことをあらかじめご了承ください。)

この記事で登場する整数は、全て素数 $m$ で割ったあまりとします(これは嘘です。指数部分などは割ってはいけません。式の結果だけ $m$ で割ってください。)

modint

ある正整数 $m$ に対して$\mod{m}$ で問題に答える場合のために、mod library を作りましょう。

modpow

整数 $a,b$ に対して $a^b \mod{m}$ の計算を行う modpow 関数を書きます。$b = \sum x_i$ のようにいくつかの整数 $x_1,x_2,...$ の和で分割したとき、以下が成り立つことに注目して計算量を減らします。

• $a^b = \Pi a^{x_i}$

$a^{x_i}$ の計算を高速に行うことができるような $x_1,x_2 ,...$ のインスタンスがどんなものかを考えます。結論から言うと、$x_1,x_2,...$ は $2$ の冪乗の形の整数を採用します。$b$ を二進数表示にすると、二冪の整数の和で $b$ を作ることができるとわかります。

$b = \sum x_i$ は $O(\log{b})$ 個の整数の和であり、$x_i \leq \log{b}$ です。

$x_i$ は $2$ の冪乗なので、$a^b$ を求めるためには $a^1,a^2,a^4,...,a^{2^{\log{b}}} $ の値がわかればよく、$i$ が小さい順に $a^{2^i} = a^{2^{i-1}} \cdot a^{2^{i-1}} $ のように計算していくことができます。

modinv

$A$ の逆数 $A^{-1}$ の計算を行う modinv 関数を書きます。 フェルマーの小定理より、$A^{m-1} = 1 \mod m $ であることを利用します。

式の両辺に $A^{-1}$ をかけることで $A^{-1} = A^{m-2} \mod m $ が得られます。なお、この式変形は $A$ と $m$ が互いに素の場合に限ります。

modcomb

$n \mathrm{C} r = \frac{n!}{r!(n-r)!} = \frac{n\mathrm{P}r}{r\mathrm{P}r}$ を計算します。$r$ が小さければそのまま $O(r+\log{m})$ 時間で計算すれば良いです。(TODO : 階乗を高速化するテクニック)

あり得る $n$ の上限 $N$ がわかっているなら、$N$ 以下の整数 $a$ に対して $a! , (a!)^{-1}$ を管理するテーブルを持っておくことで $n \mathrm{C} r = (n!\cdot (r!)^{-1})\cdot ((n-r)!)^{-1} $ をクエリあたり $O(1)$ 時間で計算できます。

FastPowerSum

整数 $n,k$ に対して、$n$ 以下の自然数の $k$ 乗和 $\sum_{1 \leq x \leq n}{x^k}$ を (ある程度) 高速に求めます。(ちなみに、過去に名古屋大学の入試数学にこれから述べる式変形を要求する問題が出ました。)

突飛ですが、以下が成り立ちます。

• $(x+1)^K = \sum_{0 \leq i \leq K}{_K\mathrm{C}_i \cdot x^i} $

これを、こうです。

• $\sum_{1\leq x \leq n}{(x+1)^K} = \sum_{0 \leq i \leq K}{(_K\mathrm{C}_i\sum_{1\leq x \leq n}{x^i})} $

右辺の $1$ つ目の $\sum$ から、$i = K-1,i=K$ の $2$ つの項を取り出すと、右辺は以下のようになります。

• $\sum_{0 \leq i \leq K-2}{(_K\mathrm{C}_i\sum_{1\leq x \leq n}{x^i})} + K \cdot \sum_{1\leq x \leq n}{x^{K-1}} + \sum_{1\leq x \leq n}{x^K}$

左辺の $\sum$ の中身を展開すると、$\sum_{1\leq x \leq n}{(x+1)^K} = (N+1)^K - 1 + \sum_{1\leq x \leq n}{x^K}$ となります。両辺から $\sum_{1\leq x \leq n}{x^K} $ を引くことで、以下の式が得られます。

• $ (N+1)^K - 1 = \sum_{0 \leq i \leq K-2}{(_K\mathrm{C}_i\sum_{1\leq x \leq n}{x^i})} + K \cdot \sum_{1\leq x \leq n}{x^{K-1}} $

これを式変形して、以下が得られます。

• $ \sum_{1\leq x \leq n}{x^{K-1}} = \frac{(N+1)^K - 1 - \sum_{0 \leq i \leq K-2}{(_K\mathrm{C}_i\sum_{1\leq x \leq n}{x^i})}} {K}$

$K$ に $k+1$ を代入することで、$ \sum_{1\leq x \leq n}{x^{k}} = \frac{(N+1)^{k+1} - 1 - \sum_{0 \leq i \leq k-1}{(_{k+1}\mathrm{C}_i\sum_{1\leq x \leq n}{x^i})}}{k+1}$ が計算できました。

右辺の $\sum_{1\leq x \leq n}{x^i}$ の結果は、$i$ が $k$ 未満なので再帰的に冪乗和の関数を呼び出すことで計算ができます。冪乗和の計算結果をメモ化しておくことで、$\sum_{1\leq x \leq n}{x^k}$ の値を先の modcomb などと組み合わせて $O(k^2)$ 時間で計算できます。

ひとまずおわり

まだまだやることはたくさんです。次回に続きます。