問題
多項式 $f(x) = \sum a_i x^i$(次数 $N-1$)と $g(x) = \sum b_j x^j$(次数 $M-1$)の積 $h(x) = f(x) g(x)$ の係数列を $\bmod 998244353$ で求めよ。
$998244353 = 119 \times 2^{23} + 1$ は NTT 素数。原始根は 3。
制約
$1 \le N, M \le 5 \times 10^5$
$0 \le a_i, b_j < 998244353$
入出力例
入力例 1
3 3
1 2 3
4 5 6
出力例 1
4 13 28 27 18
$(1+2x+3x^2)(4+5x+6x^2) = 4 + 13x + 28x^2 + 27x^3 + 18x^4$
ヒント (段階的開示)
ヒント1: 方向性
NTT は FFT の整数 mod 版。NTT 素数を使えば浮動誤差なしで $O((N+M)\log(N+M))$ の多項式乗算が可能。
ヒント2: アプローチ
①長さを $2^k \ge N+M-1$ にゼロパディング ②NTT で点値に変換 ③点ごとに積 ④INTT で係数に戻す。単位根 $\omega = g^{(p-1)/n} \bmod p$。
ヒント3: 誘導
# Cooley-Tukey バタフライ
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
for k in range(length // 2):
u = a[i+k]; v = a[i+k+length//2] * wn % MOD
a[i+k] = (u + v) % MOD
a[i+k+length//2] = (u - v) % MOD
wn = wn * w % MOD
length <<= 1
模範解答 (Python)
import sys
input = sys.stdin.readline
def solve():
MOD = 998244353
G_ROOT = 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_ROOT, (MOD - 1) // length, MOD)
if invert:
w = pow(w, MOD - 2, MOD)
for i in range(0, n, length):
wn = 1
for k in range(length // 2):
u = a[i + k]
v = a[i + k + length // 2] * wn % MOD
a[i + k] = (u + v) % MOD
a[i + k + length // 2] = (u - v) % MOD
wn = wn * w % MOD
length <<= 1
if invert:
inv_n = pow(n, MOD - 2, MOD)
for i in range(n):
a[i] = a[i] * inv_n % MOD
def poly_mul(a, b):
result_len = len(a) + len(b) - 1
n = 1
while n < result_len:
n <<= 1
fa = a + [0] * (n - len(a))
fb = b + [0] * (n - len(b))
ntt(fa, False)
ntt(fb, False)
for i in range(n):
fa[i] = fa[i] * fb[i] % MOD
ntt(fa, True)
return fa[:result_len]
N, M = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
print(*poly_mul(a, b))
solve()
Step-by-Step 解説
1NTT の基本
$\omega_n \equiv g^{(p-1)/n} \pmod p$ を使い、FFT を整数体上で実行。$998244353 - 1 = 119 \times 2^{23}$ なので $n \le 2^{23}$ まで OK。
$\omega_n \equiv g^{(p-1)/n} \pmod p$ を使い、FFT を整数体上で実行。$998244353 - 1 = 119 \times 2^{23}$ なので $n \le 2^{23}$ まで OK。
2Cooley-Tukey
DFT を再帰分割し、ビット反転置換でインプレース化。計算量 $O(n \log n)$。
DFT を再帰分割し、ビット反転置換でインプレース化。計算量 $O(n \log n)$。
3逆 NTT
$\omega$ を $\omega^{-1}$ に替えて NTT し、最後に $n^{-1}$ を掛ける。
$\omega$ を $\omega^{-1}$ に替えて NTT し、最後に $n^{-1}$ を掛ける。
4全体フロー
パディング → NTT($f$), NTT($g$) → 点積 → INTT。$O((N+M)\log(N+M))$。
パディング → NTT($f$), NTT($g$) → 点積 → INTT。$O((N+M)\log(N+M))$。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| NTT 素数でない mod を使う | 2の冪乗根が存在しない | 998244353 や 469762049 を使う |
| 長さを 2 の冪にしない | NTT の前提条件 | while n < result_len: n <<= 1 |
| 逆変換で $1/n$ 倍を忘れる | INTT の正規化 | a[i] = a[i] * inv_n % MOD |
次のステップ
- 形式的冪級数の逆元・平方根・log・exp
- 任意 mod での畳み込み(3-NTT)
- Karatsuba 法との比較