問題
次数 $N-1$ の多項式 $f(x) = \sum_{i=0}^{N-1} a_i x^i$ と次数 $M-1$ の多項式 $g(x) = \sum_{j=0}^{M-1} b_j x^j$ が与えられる（係数は $\bmod p$、$p = 998244353$）。
$h(x) = f(g(x)) \bmod x^N$ を計算し、係数 $h_0, h_1, \ldots, h_{N-1}$ を出力せよ。
制約
| パラメータ | 範囲 |
|---|---|
| $N, M$（部分点） | $1 \le N, M \le 2000$ |
| $N, M$（満点） | $N, M \le 3 \times 10^4$ |
| 係数 | $0 \le a_i, b_j < p = 998244353$ |
| 時間制限 | 4秒 |
入出力例
入力例 1
4 3
1 2 3 4
0 1 1出力例 1
1 2 5 16$f(x)=1+2x+3x^2+4x^3$、$g(x)=x+x^2$。$f(g(x)) = 1 + 2(x+x^2) + 3(x+x^2)^2 + 4(x+x^2)^3 \bmod x^4 = 1+2x+5x^2+16x^3$。
概念図: Baby-Step Giant-Step の分解
ヒント（段階的開示）
ヒント1: 方向性
愚直に $g(x)^k$ を順次計算すると $O(N^2 M)$ で TLE。$K = \lceil\sqrt{N}\rceil$ として $g^k = (g^K)^{j} \cdot g^{i}$（$k = jK + i$）と分解する Baby-Step Giant-Step アプローチで $O(N^{1.5} \log N)$。
ヒント2: アプローチ
$$f(g) = \sum_{j=0}^{K} G^j \cdot \underbrace{\left(\sum_{i=0}^{K-1} a_{jK+i} \cdot g^i\right)}_{\text{内側の和（多項式）}}$$
- Baby Steps: $g^0, g^1, \ldots, g^{K-1} \bmod x^N$ を NTT で前計算
- $G = g^K \bmod x^N$ を計算
- Giant Steps: $j = 0, 1, \ldots, K$ のループで、内側の和を計算し $G^j$ を掛けて result に加算
ヒント3: NTT の実装と全体骨格
K = int(N**0.5) + 1
g_baby = [[0]*N for _ in range(K)]
g_baby[0][0] = 1 # g^0 = 1
g_trunc = b[:N] + [0]*max(0, N-len(b))
for i in range(1, K):
g_baby[i] = poly_mul_mod(g_baby[i-1], g_trunc, N)
G = poly_mul_mod(g_baby[K-1], g_trunc, N) # g^K
result = [0]*N
G_pow = [0]*N; G_pow[0] = 1 # G^0
for j in range(K+2):
coeff_sum = [0]*N
for i in range(K):
idx = j*K + i
if idx < N and a[idx]:
for pos in range(N):
coeff_sum[pos] = (coeff_sum[pos] + a[idx]*g_baby[i][pos]) % MOD
if any(coeff_sum):
term = poly_mul_mod(G_pow, coeff_sum, N)
for pos in range(N): result[pos] = (result[pos]+term[pos]) % MOD
if j*K >= N: break
G_pow = poly_mul_mod(G_pow, G, N)模範解答 (Python)
import sys
input = sys.stdin.readline
MOD = 998244353; g_pr = 3
def ntt(a, inv=False):
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=N: break
ai=a[idx]
if ai==0: continue
for pos in range(N): cs[pos]=(cs[pos]+ai*gb[i][pos])%MOD
if any(cs):
term=pmm(Gp,cs,N)
for pos in range(N): res[pos]=(res[pos]+term[pos])%MOD
if j*K>=N: break
Gp=pmm(Gp,G,N)
print(*res)
solve()
Step-by-Step 解説
1問題の分解（BSGS）
$f(g) = \sum_{k} a_k g^k$ を $k = jK + i$（$0 \le i < K$）で分解：
$$f(g) = \sum_{j=0}^{K} G^j \cdot \left(\sum_{i=0}^{K-1} a_{jK+i} g^i\right)$$ ここで $G = g^K$。
$f(g) = \sum_{k} a_k g^k$ を $k = jK + i$（$0 \le i < K$）で分解：
$$f(g) = \sum_{j=0}^{K} G^j \cdot \left(\sum_{i=0}^{K-1} a_{jK+i} g^i\right)$$ ここで $G = g^K$。
2Baby Steps の前計算
$g^0, g^1, \ldots, g^{K-1}$ を順次 NTT 乗算で計算し $\bmod x^N$ で切り捨て。各 $O(N \log N)$、合計 $O(K N \log N) = O(N^{1.5} \log N)$。
$g^0, g^1, \ldots, g^{K-1}$ を順次 NTT 乗算で計算し $\bmod x^N$ で切り捨て。各 $O(N \log N)$、合計 $O(K N \log N) = O(N^{1.5} \log N)$。
3Giant Steps の計算
$j$ のループ内で「内側の和」= $\sum_i a_{jK+i} g^i$ を係数の線形結合として計算（$O(KN)$）。$G^j$ を掛けて result に加算。
$j$ のループ内で「内側の和」= $\sum_i a_{jK+i} g^i$ を係数の線形結合として計算（$O(KN)$）。$G^j$ を掛けて result に加算。
4NTT（数論変換）
$p = 998244353 = 119 \times 2^{23} + 1$ は NTT 素数。原始根は 3。ビット逆順転置 + butterfly で $O(N \log N)$。
$p = 998244353 = 119 \times 2^{23} + 1$ は NTT 素数。原始根は 3。ビット逆順転置 + butterfly で $O(N \log N)$。
計算量
Baby Steps 前計算: $O(K \cdot N \log N) = O(N^{1.5} \log N)$
Giant Steps: $O(K \cdot N \log N) = O(N^{1.5} \log N)$
内側の和の計算: $O(K^2 N) = O(N^2)$（ボトルネックに注意）
全体: $O(N^{1.5} \log N)$ または $O(N^2)$ （内側の愚直部分）
空間: $O(K \cdot N) = O(N^{1.5})$
Giant Steps: $O(K \cdot N \log N) = O(N^{1.5} \log N)$
内側の和の計算: $O(K^2 N) = O(N^2)$（ボトルネックに注意）
全体: $O(N^{1.5} \log N)$ または $O(N^2)$ （内側の愚直部分）
空間: $O(K \cdot N) = O(N^{1.5})$
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
poly_mul_mod で $x^N$ 以上の係数を残す | メモリ・精度の無駄 | r[:size] で切り捨て |
| $K$ の設定が小さすぎる | $j$ のループが足りず未計算の項が残る | K = int(N**0.5) + 1 + ループを $K+2$ まで |
g_baby[0][0] = 1 の初期化忘れ | $g^0 = 1$ でなく全ゼロになる | 明示的に設定 |
| Giant Step で $G$ でなく $g$ を使う | $g^{jK}$ でなく $g^j$ になってしまう | G = g^K を別途計算 |
次のステップ
- 発展問題: $f(g(x)) \bmod x^N$ を $O(N^{1.5})$ でできる特殊ケース（$g(0)=0$）
- 関連: Bostan-Mori アルゴリズムと多項式合成の組み合わせ
- 応用: 行列値多項式の合成（行列累乗との融合）