問題
$\bmod{998244353}$ で計算する。長さ$N$の数列$c_1,\dots,c_N$と$f_1,\dots,f_N$が与えられる。$a_0=1$とし、$i=1,\dots,N$について
$$a_i = c_i + \sum_{j=1}^{i} f_j \cdot a_{i-j} \pmod{998244353}$$
で定義される数列$a_1,\dots,a_N$を求めよ。
入力形式
N
c_1 c_2 ... c_N
f_1 f_2 ... f_N制約
$1 \le N \le 200000$
$0 \le c_i, f_i < 998244353$
入出力例
入力例1
3
5 3 7
2 1 4出力例1
7 18 54a0=1。a1=5+2·1=7。a2=3+2·7+1·1=18。a3=7+2·18+1·7+4·1=54。
概念図
ヒント（段階的開示）
ヒント1（方向性）
$a_i$の定義に過去の$a$との畳み込みが入っている「オンライン畳み込み」。愚直だと$O(N^2)$。畳み込みなのでNTTが使えそうだが、まだ確定していない未来の$a$を使わずに畳み込む工夫が必要。
ヒント2（アプローチ）
区間$[l,r)$の分割統治`solve(l,r)`を考える。`solve(l,mid)`で左半分を確定させ、`a[l:mid]`と`f[1:r-l]`をNTTで畳み込んだ結果を`a[mid:r]`に一括加算してから`solve(mid,r)`を呼ぶ。
ヒント3（誘導）
def solve(l, r, f, a):
if r - l == 1:
return
mid = (l + r) // 2
solve(l, mid, f, a)
left = a[l:mid]
fpart = f[1:r - l]
if fpart and left:
conv = multiply(left, fpart)
for i in range(mid, r):
idx = i - l - 1
if 0 <= idx < len(conv):
a[i] = (a[i] + conv[idx]) % MOD
solve(mid, r, f, a)a[0]=1, a[i]=c_i(i≥1)としてsolve(0, N+1, f, a)を呼べばa[1:]が答え。
模範解答 (Python)
import sys
MOD = 998244353
G = 3
def ntt(a, invert):
n = len(a)
j = 0
for i in range(1, n):
bit = n >> 1
while j & bit:
j ^= bit
bit >>= 1
j |= bit
if i < j:
a[i], a[j] = a[j], a[i]
length = 2
while length <= n:
w = pow(G, (MOD - 1) // length, MOD)
if invert:
w = pow(w, MOD - 2, MOD)
for i in range(0, n, length):
wn = 1
half = length // 2
for k in range(half):
u = a[i + k]
v = a[i + k + half] * wn % MOD
a[i + k] = (u + v) % MOD
a[i + k + half] = (u - v) % MOD
wn = wn * w % MOD
length <<= 1
if invert:
n_inv = pow(n, MOD - 2, MOD)
for i in range(n):
a[i] = a[i] * n_inv % MOD
def multiply(a, b):
result_len = len(a) + len(b) - 1
sz = 1
while sz < result_len:
sz <<= 1
fa = a + [0] * (sz - len(a))
fb = b + [0] * (sz - len(b))
ntt(fa, False)
ntt(fb, False)
for i in range(sz):
fa[i] = fa[i] * fb[i] % MOD
ntt(fa, True)
return fa[:result_len]
def solve(l, r, f, a):
if r - l == 1:
return
mid = (l + r) // 2
solve(l, mid, f, a)
left = a[l:mid]
fpart = f[1:r - l]
if fpart and left:
conv = multiply(left, fpart)
for i in range(mid, r):
idx = i - l - 1
if 0 <= idx < len(conv):
a[i] = (a[i] + conv[idx]) % MOD
solve(mid, r, f, a)
def main():
data = sys.stdin.read().split()
idx = 0
n = int(data[idx]); idx += 1
c = [0] + [int(data[idx + i]) for i in range(n)]; idx += n
f = [0] + [int(data[idx + i]) for i in range(n)]; idx += n
a = [0] * (n + 1)
a[0] = 1
for i in range(1, n + 1):
a[i] = c[i]
solve(0, n + 1, f, a)
print(' '.join(map(str, a[1:])))
main()Step-by-Step 解説
1愚直畳み込みが使えない理由
$a_i$の定義右辺に未確定に見える$a_{i-1},\dots,a_0$が現れ、一括NTT畳み込みをすると循環してしまう。
$a_i$の定義右辺に未確定に見える$a_{i-1},\dots,a_0$が現れ、一括NTT畳み込みをすると循環してしまう。
2時間軸の分離
solve(l,r)は「前半[l,mid)が後半[mid,r)に与える影響」だけを切り出す。前半はsolve(l,mid)終了時点で確定済みなので安全にNTTできる。
solve(l,r)は「前半[l,mid)が後半[mid,r)に与える影響」だけを切り出す。前半はsolve(l,mid)終了時点で確定済みなので安全にNTTできる。
3畳み込む範囲
a[l:mid]（長さmid-l）とf[1:r-l]（長さr-l-1）の畳み込みの添字kの項はΣa_{l+p}f_{k-p}となり、a[mid+idx]に正しく加算される。
a[l:mid]（長さmid-l）とf[1:r-l]（長さr-l-1）の畳み込みの添字kの項はΣa_{l+p}f_{k-p}となり、a[mid+idx]に正しく加算される。
4計算量解析
深さdでは区間長N/2^dの呼び出しが2^d個あり、各畳み込みO((N/2^d)log(N/2^d))。全体でO(N log²N)。
深さdでは区間長N/2^dの呼び出しが2^d個あり、各畳み込みO((N/2^d)log(N/2^d))。全体でO(N log²N)。
5実装上の注意
fpartやleftが空のケースをスキップしないとNTTのサイズ計算でバグを踏む。
fpartやleftが空のケースをスキップしないとNTTのサイズ計算でバグを踏む。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| a0の初期値を1にし忘れる | a0は数列に含まれないと錯覚 | a[0]=1を明示してからsolve(0,N+1,...)を呼ぶ |
| solve(mid,r)を先に呼んでしまう | 影響の向きを逆にする | 「左確定→畳み込みで右に反映→右再帰」の順を守る |
| f[1:r-l]の範囲をずらす | オフバイワン | 畳み込みたい項数は区間長r-l分なのでf[1:r-l]が正しい |
| NTTのサイズを結果長ちょうどにする | 2冪でないと基底変換できない | szを結果長以上の最小の2冪に丸める |
次のステップ
- 発展: この技法は多項式のexp・log・逆元計算のNewton法の中の「オンライン乗算」の基礎になっている
- 次回予告: B-Tree（2-3-4木の挿入とノード分裂）