Day 081-Q4 — 多次元 NTT（テンソル積畳み込み・2D 多項式積）

問題

2 次元多項式 $F(x, y) = \sum_{i=0}^{N-1}\sum_{j=0}^{M-1} a_{i,j} x^i y^j$ と $G(x, y)$ の積 $H = F \cdot G$ を求めよ。

$H(x, y) = \sum_{k=0}^{2N-2}\sum_{l=0}^{2M-2} c_{k,l} x^k y^l$ の全係数 $c_{k,l}$ を $998244353$ で割った余りで出力せよ。

入出力例

2 2
1 2
3 4
5 6
7 8

5 16 12
22 60 40
21 52 32

概念図: 2D-NTT の処理フロー

ヒント

def ntt2d(A, H, W, invert):
    # 行方向 NTT
    for i in range(H): ntt(A[i], invert)
    # 列方向 NTT
    for j in range(W):
        col = [A[i][j] for i in range(H)]
        ntt(col, invert)
        for i in range(H): A[i][j] = col[i]

模範解答

import sys
input = sys.stdin.readline

MOD = 998244353
g_prim = 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_prim, (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 ntt2d(A, H, W, invert):
    for i in range(H): ntt(A[i], invert)
    for j in range(W):
        col = [A[i][j] for i in range(H)]
        ntt(col, invert)
        for i in range(H): A[i][j] = col[i]

def solve():
    N, M = map(int, input().split())
    A = [list(map(int, input().split())) for _ in range(N)]
    B = [list(map(int, input().split())) for _ in range(N)]

    RH = 1
    while RH < 2*N-1: RH <<= 1
    RW = 1
    while RW < 2*M-1: RW <<= 1

    FA = [[0]*RW for _ in range(RH)]
    FB = [[0]*RW for _ in range(RH)]
    for i in range(N):
        for j in range(M):
            FA[i][j] = A[i][j]; FB[i][j] = B[i][j]

    ntt2d(FA, RH, RW, False)
    ntt2d(FB, RH, RW, False)

    FC = [[FA[i][j]*FB[i][j]%MOD for j in range(RW)] for i in range(RH)]
    ntt2d(FC, RH, RW, True)

    out = []
    for i in range(2*N-1):
        out.append(' '.join(str(FC[i][j]) for j in range(2*M-1)))
    print('\n'.join(out))

solve()

Step-by-Step 解説

Step 1: 2次元多項式積の基本原理

$H[k][l] = \sum_{i+j=k}\sum_{p+q=l} F[i][p] \cdot G[j][q]$ — 「行方向の畳み込み」と「列方向の畳み込み」を独立に適用することと等価。

Step 2: テンソル積としての理解

2D-NTT は $\text{NTT}_H \otimes \text{NTT}_W$（Kronecker 積）。行→列の順でも列→行の順でも同じ変換結果になる（可換性）。

Step 3: ゼロパディングサイズの決定

畳み込み結果のサイズ $(2N-1) \times (2M-1)$ なので，$\ge 2N-1$ の最小 2 の冪 $RH$，$\ge 2M-1$ の最小 2 の冪 $RW$ が必要。

Step 4: 計算量の確認

行方向 NTT: $O(RH \cdot RW \log RW)$，列方向 NTT: $O(RH \cdot RW \log RH)$，全体: $O(NM \log(NM))$。逆 NTT も同様。

よくあるミス

次のステップ

• 発展問題: 3次元 NTT（$N \times M \times L$ 多項式の積）
• 多変数多項式剰余（mod $p(x), q(y)$）の高速計算

自己評価