1. はじめに
2. 観察
3. \(K\)乗和の公式
4. \(f(N, M)\)の性質
5. ラグランジュ補間
6. ラグランジュ補間を利用した\(f(N)\)の評価
7. まとめ
8. Bonus

1.　はじめに

この問題は、高難度の数え上げでしばしば部分問題として出てくるテクニックを取り上げた問題です。この問題を解くことで得られる知見が、将来より高度な数え上げ問題を AC するための足掛かりになったら幸いです。

この問題は数列の第 \(n\) 項を高速に計算する問題です。この問題のように線形よりも高速に計算する必要がある問題では、高速化のために数列の性質を理解することが重要になります。以下の表では、 ABC で過去に多く出題されている数列の性質・および計算量をまとめたものです。

数列の種類	具体例	前計算	クエリ
累乗	\(n ^ k\)	なし	\(\mathrm{O}(\log k)\)
階乗・二項係数	\(n!,\binom{n}{k}\)	\(\mathrm{O}(n)\)	\(\mathrm{O}(1)\)
遷移が\(d \times d\)の　 　行列で表せる数列	フィボナッチ数列	なし	\(\mathrm{O}(d^3 \log k)\)

さて、この問題では題意の数列はどのような性質を満たすのでしょうか？

2.　観察

このような問題で行き詰まった時は、簡単なケースでの状態を観察してよい性質を見つけるというのが考察の方法論の 1 つとして考えられます。

実際に \(m\) が小さいケースについて考えてみましょう。まず初めに、 \(m = 0\) の時は簡単で

\[ f(n, 1) = n^k\]

であるから容易に計算できます。次に \(m = 1\) の場合を考えてみると、

\[ \begin{aligned} f(1, 1) &= 1^K \\ f(2, 1) &= 1^K + 2^K \\ f(3, 1) &= 1^K + 2^K + 3^K \\ \vdots \end{aligned} \]

となることから、 \(f(n, 1)\) (以下、便宜上 \(S_K(n)\) と書きます)は \(K\) 乗和、すなわち

\[ S_K(n) = \sum_{i = 1} ^ n i^K \]

であることがわかります。

特に \(K = 1, 2, 3\) の場合は高校の授業で出てくるなじみ深い公式が \(S_K(n)\) の値になり、答えは次の通りになります。

• \(K = 1\) のとき \(S_1(n) = \frac{n(n + 1)}{2} \)
• \(K = 2\) のとき \(S_2(n) = \frac{n(n+1)(2n+1)}{6} \)
• \(K = 3\) のとき \(S_3(n) = \frac{n^2 (n+1)^2}{4} \)

ここで上の式を観察すると、\(S_K(n)\) は \(K + 1\) 次式であるという重要な仮説が立ちます。この仮説を次の章で証明していきましょう。

3. \(K\) 乗和の公式

次の命題を証明します。

\(\displaystyle S_K(N) = \sum_{n=1}^N n^k\)は \(K + 1\) 次式で表せる。

証明はいくつかの方法がありますが、ここでは帰納法を利用した証明を取り上げます。

関数

\[T_K(n)=n^K - (n-1)^K\]

の累積和

\[ \sum_{n=1}^N T_{K+1}(n) \]

を二つの方法で表すことを考えます。まず、

\[ \begin{aligned} \displaystyle \sum_{n=1}^N T_{K+1}(n) &= \sum_{n=1}^N (n+1)^K - n^K \\ &= (N^1 - N^0) + \dots + (N^{K+1} - N^K) \\ &= N^{K+1} \end{aligned} \]

です。次に、

\[ \begin{aligned} T_{K+1}(n) &= n^{K+1} - \sum_{i=0}^{K+1} \binom{K+1}{i} n^i (-1)^{K+1-i} \\ &= (K + 1)n^K + ( nの K - 1 次式 ) \end{aligned} \]

なのを利用すると、

\[\displaystyle \sum_{n=1}^N T_{K+1}(n) = \sum_{n=1}^N (K + 1)n^K + (nのK-1次式)\]

とも表せます。 \(K - 1\) 以下の数に対して命題が成り立つとき、 \( (nのK-1次式)\) の\(1\)から\(N\)までの累積和は \((NのK次式)\) になるのを利用すると、上式は

\[\sum_{n=1}^N T_{K+1}(n) = (NのK次式) + (K + 1) \sum_{n=1}^N n^K\]

になります。以上より得られた二つの式を連立すると

\[S_K(N) = \sum_{n=1}^N n^K = \frac{N^{K+1} + (NのK次式)}{K + 1}\]

が従います。よって \(K=1\) の場合と合わせて帰納的に示せました。

4. \(f(N, M)\)の性質

元の問題に戻りましょう。 3 で示された事実から、 \(f(N,M)\) は次の性質を満たすことが証明できます。

\(f(N,M)\) は \(M\) を固定したとき \(N\) の \(M + K\) 次式で表せる。

この事実は帰納的に示せます。\(f(N,m)\) が \(N\) の \(m+K\) 次式で表せるとき、3 で得た事実と合わせると

\[ \begin{aligned} f(N, m+1) &= \sum_{n = 1} ^ N f(n, m) \\ &= \sum_{n=1}^N (nのm+K次式) \\ &= (Nのm+K+1次式) \end{aligned} \]

が従うので、 \(m=0\) のとき \(f(n,0)=n^K\) と合わせて証明できました。

以上より求める数列は、\(M\) を固定すれば \(N\) の多項式で表せるとわかりました。しかしここで新たな問題が発生します。それは「 \(f(n)\) が多項式であると分かったときに \(f(N)\) をどのように高速に計算するのか？」という問題です。

もしも多項式の係数が分かっている場合はホーナー法などのナイーブなアルゴリズムで \(\mathrm{O}(N)\) で計算できます。しかしこの問題では \(f(N, M)\) の係数を復元するのは難しそうに見えます。

次の章以降では、多項式の持つ性質をいくつか挙げた上で \(f(N)\) を高速に計算するアルゴリズムを説明していきます。

5. ラグランジュ補間

多項式は実は次の性質を持っています。

\(x_0, x_1, \dots,x_d\) と \(y_0, \dots,y_d\) が与えられたとき、

\[f(x_i) = y_i (i=0,\ldots,d)\]

を満たす \(d\) 次以下の多項式 \(f(x)\) はただ1つに定まる。(ただし \(x_i\) は互いに異なるとする。)

一次式を例に考えると、「 \(f(0) = b, f(1) = a + b\) を満たす \(1\) 次以下の式」は \(f(x) = ax + b\) に限られることがわかります。これを一般に拡張したのが上の命題です。

この命題は次の重要な事実を導きます。

• 係数が未知である \(d\) 次の多項式 \(f(n)\) が与えられたとき、\(f(N)\) を計算したい。このとき、多項式の係数がわからなくても \(f(0), f(1), \dots, f(d)\) が計算できれば \(f(N)\) を復元することが出来る。

元の問題を振り返ると、 \(f(0, M), f(1, M), \dots, f(d, M)\) は累積和を利用すれば容易に計算できる問題でした。このような問題では、多項式を \(f(0), f(1),\dots, f(d)\) の形で持つことがアルゴリズムの計算量削減に大きく役立ちます。

さて、この命題は行列を利用して証明できます。

\[f(x) = f_0 + f_1 x + f_2 x^2 + \dots + f_d x^d\]

とおくと、\((f_0, f_1,\dots,f_d)\) と \((y_0,\dots,y_d)\) はヴァンデルモンド行列を利用して

\[ \left( \begin{array}{ccccc} 1 & x_0 & x_0^2 & \cdots & x_0^d \\ 1 & x_1 & x_1^2 & \cdots & x_1^d \\ 1 & x_2 & x_2^2 & \cdots & x_2^d \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & x_d & x_d^2 & \cdots & x_d^d \\ \end{array} \right) \left( \begin{array}{c} f_0 \\ f_1 \\ f_2 \\ \vdots \\ f_d \\ \end{array} \right) = \left( \begin{array}{c} y_0 \\ y_1 \\ y_2 \\ \vdots \\ y_d \\ \end{array} \right) \]

と表せます。ヴァンデルモンド行列の行列式は

\[ \left \vert \begin{array}{ccccc} 1 & x_0 & x_0^2 & \cdots & x_0^d \\ 1 & x_1 & x_1^2 & \cdots & x_1^d \\ 1 & x_2 & x_2^2 & \cdots & x_2^d \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & x_d & x_d^2 & \cdots & x_d^d \\ \end{array} \right \vert = \prod_{j \lt i} (x_i - x_j) \neq 0 \]

であるので、条件を満たす \(f_0,\ldots,f_d\) は一意に定まることが示されました。

以上より\(f(0),f(1),\ldots,f(d)\) が与えられたときの \(d\) 次の多項式 \(f(x)\) が一意に定まること分かりましたが、ここから \(f(N)\) を求めるために \(y\) のベクトルにヴァンデルモンド行列を掛けるのは計算量が重いので、ラグランジュ基底多項式と呼ばれる多項式を導入します。

ラグランジュ基底多項式 \(L_i (x)\) を \(x_0,\dots,x_d\) を用いて次のように定めます。

\[ \begin{aligned} L_i(x) &= \prod_{j \neq i} \frac{x - x_j}{x_i - x_j} \\ &= \frac{(x - x_0) (x - x_1)\dots(x-x_{i-1})(x-x_{i+1})\dots (x-x_d)}{(x_i - x_0) (x_i - x_1)\dots(x_i-x_{i-1})(x_i-x_{i+1})\dots (x_i-x_d)} \end{aligned} \]

このとき、 \(f(x)\)は\(L_i(x)\) を用いて次のように表すことが出来ます。

\(x_0, x_1, \dots,x_d\)と\(y_0, \dots,y_d\)が与えられたとき、

\[f(x_i) = y_i (i=0,\ldots,d)\]

を満たす \(d\) 次以下の多項式 \(f(x)\) はラグランジュ基底多項式\(L_i(x)\)を用いて次のように表せる。

\[f(x) = y_0 L_0(x) + y_1 L_1(x) + \dots + y_d L_d(x)\]

上式の正当性を確かめます。ラグランジュ基底多項式

\[ L_i(x) = \frac{(x - x_0) (x - x_1)\dots(x-x_{i-1})(x-x_{i+1})\dots (x-x_d)}{(x_i - x_0) (x_i - x_1)\dots(x_i-x_{i-1})(x_i-x_{i+1})\dots (x_i-x_d)} \]

は次の性質を満たします。

\[L_i(x_j) = \begin{cases} 1 & i = j \\ 0 & i \neq j \end{cases} \]

上式より

\[ y_0 L_0(x_i) + y_1 L_1(x_i) + \dots + y_d L_d(x_i) = y_i\]

が従い、 \(L_i(x)\) が高々 \(d\) 次なのと合わせて正当性が確かめられます。

以上より、点 \((x, y)\) が複数与えられたときに、与えられた点全てを通る次数最小の多項式は一意に定まることが示されました。このように与えられた点を通る多項式を発見することを補間と呼び、特にラグランジュ基底多項式を用いた補間をラグランジュ補間と呼びます。

6. ラグランジュ補間を利用した \(f(N)\) の評価

さて、準備がようやく整いました。今までの議論を元に、\(f(0),f(1),\dots,f(d)\) から \(f(N)\) を高速に計算するアルゴリズムを説明します。

\(y_0 = f(0),y_1 = f(1) \ldots, y_d = f(d)\) が既知である \(d\) 次以下の多項式 \(f(x)\) に対して、\(f(N)\) は \(\mathrm{O}(d)\) で計算できる。

\(x_i = i\) とした場合のラグランジュ補間多項式を利用すると

\[f(N) = y_0 L_0(N) + y_1 L_1(N) + \dots + y_d L_d(N)\]

と表せるので、\(L_i(N)\) を高速に計算できればよいとわかります。\(L_i(N)\) は

\[ \begin{aligned} L_i(N) &= \frac {(N - 0) (N - 1)\dots(N-(i-1))(N-(i+1))\dots (N-d)} {(i - 0) (i - 1)\dots(i-(i-1))(i-(i+1))\dots (i-d)} \\ &= \frac {\displaystyle \prod_{0 \leq j \leq d, j \neq i} (N - j)} {i! \cdot (-1)^{d - i} (d - i)!} \end{aligned} \]

と変形できます。\(N - i\) の左右からの累積積と\((1!)^{-1},(2!)^{-1},\dots,(d!)^{-1}\) をともに \(\mathrm{O}(d)\) で前計算しておくと、\(L_i(N)\) は \(\mathrm{O}(1)\) で計算できます。よって \(L_i(N)\) を全体で \(\mathrm{O}(d)\) で計算することが出来たので、\(f(N)\) も \(\mathrm{O}(d)\) で求めることが出来ました。

7. まとめ

以上の考察により、この問題は次のアルゴリズムで \(\mathrm{O}(K \log K + (K + M)M)\) で解けることが分かります。

• \(f(n, 0) = n^K\) の初めの \(K+M+1\) 項を愚直に求める。
• 累積和を \(M\) 回計算して \(f(n, M)\) の初めの \(K+M+1\) 項を求める。
• ラグランジュ補間で \(f(N, M)\) を求める。

実装は次の通りです。TL が少し厳しいため、PyPy で AC する際には JIT コンパイルによる除算の高速化や多倍長演算によるボトルネックに配慮した実装をしないと TLE してしまうかもしれません。

mod = 1000000007

def fast_pow(x, k):
  res = 1
  while k:
    if k & 1:
      res = res * x % mod
    x = x * x % mod
    k >>= 1
  return res

N, M, K = map(int, input().split())
S = M + K
dp = [fast_pow(i, K) for i in range(S + 1)]
for _ in range(M):
  for i in range(1, S + 1):
    dp[i] = (dp[i - 1] + dp[i]) % mod

invfac = [0] * (S + 1)
fac = 1
for i in range(1, S + 1):
  fac = fac * i % mod
invfac[-1] = pow(fac, mod - 2, mod)
for i in reversed(range(S)):
  invfac[i] = invfac[i + 1] * (i + 1) % mod

N %= mod

L = [1] * (S + 1)
for i in range(S):
  L[i + 1] = L[i] * (N - i) % mod
R = [1] * (S + 1)
for i in reversed(range(1, S + 1)):
  R[i - 1] = R[i] * (N - i) % mod

answer = 0
for i in range(S + 1):
  tmp = dp[i] * L[i] % mod * R[i] % mod * invfac[i] % mod * invfac[S - i] % mod
  if S % 2 == i % 2:
    answer += tmp
  else:
    answer -= tmp

print(answer % mod)

8. Bonus

1. この問題では多項式の係数列 \((f_0, f_1,\dots,f_d)\) は必要が無かったので計算しませんでしたが、もし係数が必要な場合はヴァンデルモンド行列を利用すると \((y_0,\dots,y_d)\) から \((f_0, f_1,\dots,f_d)\) を \(\mathrm{O}(d \log^2 d)\) で計算できることが知られています。Library Checker の Polynomial Interpolation に実装例が置いてあるので興味がある方は調べてみてください。

2. \(N\leq 10^{18}, M \leq 10^7, K \leq 10^6\) とします。\(f(N, M) \bmod 998244353\) を計算してください。