目標

長さ $N$ の数列 $a_0, a_1, a_2,\cdots, a_{N-1}$ に対して累積和を $K$ 回取った数列 $b_0, b_1, b_2,\cdots, b_{N-1}$ の $b_{N-1}$ を求めてください。

$1\leq N\leq 10^6 , 1\leq K \leq 10^9, 1\leq a_i, b_i\leq 10^9$
答えを $998244353$ で割った余りを出力してください。

例えば、数列 $1,2,3$ に対して $K=2$ で操作をすれば $1,4,10$ が得られて、答えは $10$ になります。

この問題は実験や組み合わせ論的解釈で解けますが、一般的にこの 2 つの方法、特に後者は天才向けの戦い方として有名です。

(図　天才向けの方法と形式的冪級数)

これに対し、形式的冪級数を使ったやり方を使うと機械的な計算で解法を導き出すことができます。この記事の目標は形式的冪級数に入門して式変形でこの問題を解くことです。

形式的冪級数

なにこれ？

$$\sum_{i=0}^\infty a_i x^i$$

項数が有限個とは限らない多項式です。専ら無限個の項を持っているということにして扱います。
無限和特有の操作もありますが、基本的には有限項の多項式みたく扱えます。

このとき、$x$ に具体的な値は代入しませんし、収束するとか発散するとかも気にしません。これ以降はあくまで形式的にこの冪級数を扱うことを考えます。

数列との対応

数列 $a_0, a_1, a_2,\cdots$ に対して、
$$A=\sum_{i=0}^\infty a_i x^i$$
なる $A$ を対応付けます。この $A$ は数列 $a_0, a_1, a_2,\cdots$ に対する母関数と呼ばれるものです。

有限長の数列の場合、
$$a_0, a_1, a_2,\cdots, a_n$$
を次のように捉えれば同様に対応付けができます。
$$a_0, a_1, a_2,\cdots, a_n, 0, 0, \cdots$$

数列の予測をしてみよう

長さ $N$ の数列 $a_0, a_1, a_2,\cdots, a_{N-1}$ に対して累積和を $K$ 回取った数列 $b_0, b_1, b_2,\cdots, b_{N-1}$ の $b_{N-1}$ を求めてください。

$1\leq N\leq 10^6 , 1\leq K \leq 10^9, 1\leq a_i, b_i\leq 10^9$
答えを $998244353$ で割った余りを出力してください。

累積和は積として分離できる

さて、任意の $a_i$、$K$ に対して数列 $b_i$ の予測をしてみましょう！

$$A=\sum_{i=0}^\infty a_i x^i$$

に対して 1 回累積和を取ったものを $f(A)$ とすると、

$$f(A)=(1+x+x^2+\cdots)\sum_{i=0}^\infty a_i x^i$$

となります。実際に前数項を見てみましょう。多項式・形式的冪級数 $A$ の $x^i$ 項の係数を $[x^i]A$ と表すことにして以下のように表せます。

$$[1]f(A)=[1]a_0=a_0$$
$$[x]f(A)=[x](a_0\times x + a_1x)=[x] (a_0x+a_1x)=a_0+a_1 $$
$$[x^2]f(A)=[x^2](a_0\times x^2 + a_1x\times x + a_2x^2)=[x^2](a_0x^2 + a_1x^2 + a_2x^2)=a_0+a_1+a_2 $$

累積和になっていますね。
$A$ の累積和を取るには「累積和を表す級数」を $A$ に掛け算すればよいということですね。
つまり、$A$ に対して $K$ 回累積和を取ったものを $f_K(A)$ とすると、

$$f_K(A)=(1+x+x^2+\cdots)^K\sum_{i=0}^\infty a_i x^i$$

これが数列 $b_i$ の正体になります。

もう解けるな。。。

この時点で、ダブリングしながら FFT で畳み込みをやっていけば $b_i$ は時間計算量 $\mathcal{O}(N\log{N} \log{K})$ で求めることができます。（え？

累積和の累乗、それから答えの計算

$b_{N-1}$ を求めるため、今度は「累積和を表す級数」の累乗についてもう少し見てみましょう。
形式的冪級数が無限級数であることを利用して、「累積和を表す級数」は次のように表せます。

$$1+x+x^2+\cdots=\frac{1}{1-x}$$

これは両辺に $(1-x)$ を乗じてみると正しいことがわかります。

$$(左辺)(1-x)=(1+x+x^2+\cdots)(1-x)=1+x+x^2+\cdots-(x+x^2+x^3+\cdots)=1$$
$$(右辺)(1-x)=\frac{1}{1-x}(1-x)=1$$

したがって、
$$f_K(A)=\frac{1}{(1-x)^K}\sum_{i=0}^\infty a_i x^i$$

ですから、二項級数に落とし込んでみると、
$$f_K(A)=(\sum_{j=0}^\infty \binom{K+j-1}{j}x^j)(\sum_{i=0}^\infty a_i x^i)$$

となります。$[x^{N-1}]$ の項をまじめに考えてみると、

$$[x^{N-1}]f_K(A)=\binom{K-1}{0}a_{N-1}x^{N-1}+\binom{K}{1}x a_{N-2}x^{N-2}+\cdots+\binom{K+N-2}{N-1}x^{N-1}a_0$$
$$= \sum_{i=0}^{N-1}\binom{K+i-1}{i}a_{N-i-1}x^{N-1}$$

これはもう $\mathcal{O}(N+K)$ で計算できますね。

ところで、
$$\binom{K+i-1}{i}=\frac{_{K+i-1}\mathrm{P} _{i}}{i!}$$

ですから、${\lbrace \frac{1}{i!}\rbrace}_{i=0}^{N-1}$ と、${\lbrace _{K+i-1}\mathrm{P} _{i} \rbrace} _{i=0}^{N-1}$ を前もって計算しておけば $\mathcal{O}(N)$ にまでできます。この時間計算量が最良です。

おわりに

この記事のために形式的冪級数をちょっと勉強しました。名前だけ知ってる状態から分かった気になってる状態に進歩した気がします。

おまけ

愚直コードを用いた本解法の verify です。

mod = 998244353

from random import randint
def make_case():
  N, K = randint(1, 100), randint(1, 100)
  yield N, K, [randint(1, 10 ** 9) for _ in range(N)]

def naive(N, K, a):
  for _ in range(K):
    for i in range(N - 1):
      a[i + 1] += a[i]
      a[i + 1] %= mod
  return a[-1]

def solve(N, K, a):
  inv_factorial = [1] * N
  factorial = [1] * N
  
  for i in range(1, N): inv_factorial[i] = inv_factorial[i - 1] * pow(i, mod - 2, mod) % mod
  for i in range(1, N): factorial[i] = factorial[i - 1] * (K + i - 1) % mod

  res = 0
  for i in range(N):
    res += factorial[i] * inv_factorial[i] * a[N - i - 1]
    res %= mod
  return res


cases_num = 100
correct = 0
for _ in range(cases_num):
  for N, K, a in make_case():
    n, s = naive(N, K, a[: ]), solve(N, K, a)
    if n == s: correct += 1
    else:
      print("Failed with: N = {}, K = {}".format(N, K))
      print(*a)
      print("Expected:", n)
      print("Received:", s)

print("Accuracy: {} / {}".format(correct, cases_num))