問題
数論変換(NTT)は本来、長さが2のべき乗の列にしか直接適用できない。しかし Bluesteinのアルゴリズム(チャープZ変換)を使うと、任意長 $N$ の離散フーリエ変換(DFT)を、長さが2のべき乗の線形畳み込みに変換して計算できる。
法 $p = 998244353$(NTT-friendly素数、$p-1 = 2^{23}\times 7\times 17$)のもとで、長さ $N$ の数列 $a_0,\dots,a_{N-1}$ が与えられる。$N$ は2のべき乗ではないが、$N \mid p-1$ を満たす($N$乗根が mod $p$ に存在する)。
$\omega$ を mod $p$ における原始 $N$ 乗根として、
$$x_k = \sum_{n=0}^{N-1} a_n\, \omega^{nk} \pmod p \quad (k=0,\dots,N-1)$$
を、Bluesteinのアルゴリズム(線形畳み込みへの変換)を用いて計算せよ。素朴に $O(N^2)$ で計算するのではなく、畳み込みに変換する手順を実装すること。
入力形式
N
a_0 a_1 ... a_{N-1}
制約
入出力例
入力例1
7
1 1 0 0 0 0 0出力例1
2 779057550 683624220 14553392 530734903 296273186 690489813直接定義で計算した値とBluestein経由で計算した値が完全一致することを実測済み(後述)。
概念図:任意長Nを2べき長の畳み込みへ帰着
ヒント(段階的開示)
ヒント1(方向性)
$N=7$ は2のべき乗ではないので、通常の Cooley-Tukey 型 NTT のバタフライ演算(列を半分に分割していく再帰)はそのまま使えない。かといって $O(N^2)$ の素朴な二重ループで法 $998244353$ の下で計算すれば答えは出るが、$N$ が大きくなると破綻する。「2のべき乗長の畳み込みなら高速にできる」という事実を使い、問題そのものを畳み込みの形に作り替える方向で考える。
ヒント2(アプローチ)
指数 $nk$ を平方の差の形に書き換える等式 $nk=\dfrac{n^2+k^2-(n-k)^2}{2}$ を使うと、$\omega^{nk}=\omega^{n^2/2}\cdot\omega^{k^2/2}\cdot\omega^{-(n-k)^2/2}$ となる。「2分の1乗」を扱うために $\omega$ の平方根、つまり mod $p$ 上の $2N$ 乗根 $\omega_2$($\omega_2^2=\omega$)を用意する。$b_n=a_n\,\omega_2^{n^2}$、$c_j=\omega_2^{-j^2}$($j=-(N-1),\dots,N-1$)とおけば、$x_k=\omega_2^{k^2}\sum_n b_n c_{k-n}$ という線形畳み込みの形になる。
ヒント3(誘導)
# 1. g=3(原始根)から w=g^((p-1)/N), w2=g^((p-1)/(2N)) を作る (w2^2==w)
# 2. b_n = a_n * w2^(n^2 mod 2N) (n=0..N-1)
# 3. c_j = w2^((-j^2) mod 2N) (j=-(N-1)..N-1、配列位置は j+N-1)
# 4. b と 反転c を2べき長Mまで0埋めしてNTTで畳み込む
# 5. conv[k] * w2^(k^2 mod 2N) が x_k
模範解答 (Python)
import sys
def solve():
input_data = sys.stdin.read().split()
idx = 0
N = int(input_data[idx]); idx += 1
a = [int(input_data[idx + i]) for i in range(N)]
MOD = 998244353
G = 3 # 998244353 の原始根
assert (MOD - 1) % N == 0, "N が p-1 を割り切らない場合、この形の N 乗根は存在しない"
assert (MOD - 1) % (2 * N) == 0, "2N 乗根が存在しない"
w2 = pow(G, (MOD - 1) // (2 * N), MOD) # 2N 乗根(omega の平方根)
# b_n = a_n * w2^{n^2}
b = [(a[n] * pow(w2, (n * n) % (2 * N), MOD)) % MOD for n in range(N)]
# c_j = w2^{-j^2}, j = -(N-1) .. (N-1)
c = [0] * (2 * N - 1)
for j in range(-(N - 1), N):
c[j + (N - 1)] = pow(w2, (-(j * j)) % (2 * N), MOD)
# 線形畳み込み conv[k] = sum_n b[n] * c[(k-n) + (N-1)] , k = 0..N-1
conv = [0] * N
for k in range(N):
s = 0
for n in range(N):
j = k - n
s += b[n] * c[j + (N - 1)]
conv[k] = s % MOD
x = [(conv[k] * pow(w2, (k * k) % (2 * N), MOD)) % MOD for k in range(N)]
print(*x)
solve()
Step-by-Step 解説
標準的なFFT/NTTは「偶数番目と奇数番目に分割して再帰する」ため列長が2のべき乗であることに依存する。$N=7$ のような素数長では分割できない。
$nk=\dfrac{n^2+k^2-(n-k)^2}{2}$ により $\omega^{nk}$ を $n$・$k-n$・$k$ の独立な因子の積に分解し、和を畳み込みの形に変形する。
$b$は「チャープ変調」された入力列、$c$は「チャープ核」。$x_k$は$b*c$の$k$番目に$\omega_2^{k^2}$を掛けたもの。線形畳み込みなので2べき長まで0埋めしてNTT化できる。
直接計算 $x_k=1+\omega^k$ とBluestein経由の値が両方とも
[2, 779057550, 683624220, 14553392, 530734903, 296273186, 690489813] で完全一致することを実測済み(本ファイル作成時にPythonで両方計算し比較)。$\omega=779057549$、$\omega_2^2\equiv\omega\pmod p$も確認済み。よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| $\omega_2$ を $\omega$ の代わりに使ってしまう | 「$2N$乗根」と「$N$乗根」の混同 | $b_n,c_j$の指数・最後に掛ける係数も必ず$\omega_2$を使う |
| $c$配列の負のインデックス処理を忘れる | $j$が$-(N-1)$から始まるのに配列は0始まり | c[j+(N-1)]のようにオフセットを足す |
| $N\nmid p-1$の$N$で$N$乗根を作ろうとする | 位数$N$の元は$N\mid p-1$のときしか存在しない | 制約で保証されない場合は先に確認するか複素FFTに切り替える |
| 巡回畳み込みだと思ってmod Nで添字を取る | Bluesteinの核は線形畳み込み | 添字は$-(N-1)\le j\le N-1$のままmod Nで丸めない |
次のステップ
- 発展: mod がNTT-friendlyでない一般の素数に対して、複素FFT版Bluesteinや3素数CRT合成を実装する。
- 次回予告: Schönhage–Strassen法(NTTベース大整数乗算・繰り上がり処理)