問題
長さ$N$の数列$A$と長さ$M$の数列$B$、および法$\mathrm{MOD}$が与えられる。畳み込み $c_k=\sum_{i+j=k}a_ib_j \pmod{\mathrm{MOD}}$($k=0,\dots,N+M-2$)を求めよ。
$\mathrm{MOD}$はNTTに適した素数とは限らない。そこで「NTTフレンドリーな3つの素数$p_1,p_2,p_3$の下でそれぞれ畳み込みを計算し、中国剰余定理(CRT)で合成してから最後に$\mathrm{MOD}$で還元する」という定石を使う。
使用する3素数: $p_1=167772161,\ p_2=469762049,\ p_3=754974721$($p_1p_2p_3\approx5.98\times10^{25}$)
手順: ①$p_1,p_2,p_3$それぞれを法として畳み込みを計算 ②各$k$について3つの余りからGarnerのアルゴリズムで$c_k\bmod(p_1p_2p_3)$を復元 ③得られた値を$\mathrm{MOD}$で割った余りを出力
入力形式
N M
a_0 ... a_{N-1}
b_0 ... b_{M-1}
MOD
制約
$1 \le N, M \le 2000$
$0 \le a_i, b_i < 10^9$
$2 \le \mathrm{MOD} \le 10^9$
入出力例
入力例1
2 2
1 2
3 4
7
出力例1
3 3 1
$(1+2x)(3+4x)=3+10x+8x^2$。各係数を7で割った余りは $3,\ 10\bmod7=3,\ 8\bmod7=1$。
概念図
ヒント(段階的開示)
ヒント1: 方向性
NTTは「法-1が大きな2の冪で割り切れる素数」でなければ動かないため、$\mathrm{MOD}$が任意の整数だとそのままNTTが使えない。「求めたい値そのもの」と「計算に使う法」を分けて考えるのがカギ。
ヒント2: アプローチ
畳み込みの真の値(modを取る前の整数値)は一定の範囲に収まる。互いに素な複数のNTTフレンドリー素数でそれぞれ畳み込みを計算し、中国剰余定理でその真の値を復元し、最後に本当に欲しい$\mathrm{MOD}$で割り直す。3つの余りから元の値を復元する定番の方法がGarnerのアルゴリズム。
ヒント3: 誘導(コード骨格)
def garner_combine(r1, r2, r3):
p1, p2, p3 = 167772161, 469762049, 754974721
t1 = r1 % p1
t2 = ((r2 - t1) * pow(p1, p2 - 2, p2)) % p2
t3 = ((r3 - t1 - t2 * p1) * pow((p1 * p2) % p3, p3 - 2, p3)) % p3
return t1 + t2 * p1 + t3 * p1 * p2
`pow(x,p-2,p)`はフェルマーの小定理によるmod逆元。$t_1,t_2,t_3$を順に確定させ、最後に$t_1+t_2p_1+t_3p_1p_2$という混合基数表現で復元する。
模範解答 (Python)
import sys
MODS = [167772161, 469762049, 754974721]
def convolve_mod(a, b, mod):
n, m = len(a), len(b)
res = [0] * (n + m - 1)
for i in range(n):
ai = a[i]
if ai == 0:
continue
for j in range(m):
res[i + j] = (res[i + j] + ai * b[j]) % mod
return res
def garner_combine(r1, r2, r3):
p1, p2, p3 = MODS
t1 = r1 % p1
t2 = ((r2 - t1) * pow(p1, p2 - 2, p2)) % p2
t3 = ((r3 - t1 - t2 * p1) * pow((p1 * p2) % p3, p3 - 2, p3)) % p3
return t1 + t2 * p1 + t3 * p1 * p2
def solve():
data = sys.stdin.read().split()
idx = 0
n = int(data[idx]); idx += 1
m = int(data[idx]); idx += 1
A = [int(x) for x in data[idx:idx + n]]; idx += n
B = [int(x) for x in data[idx:idx + m]]; idx += m
MOD = int(data[idx]); idx += 1
c1 = convolve_mod(A, B, MODS[0])
c2 = convolve_mod(A, B, MODS[1])
c3 = convolve_mod(A, B, MODS[2])
result = []
for k in range(n + m - 1):
v = garner_combine(c1[k], c2[k], c3[k])
result.append(v % MOD)
print(*result)
solve()
計算量: $O(NM)$(素朴な畳み込み×3)+ $O(N+M)$(Garner合成)。実用では畳み込みをNTTに置き換え$O((N+M)\log(N+M))$。
Step-by-Step 解説
13つの法それぞれで畳み込み
本問の規模なら素朴な$O(NM)$畳み込みで十分。
本問の規模なら素朴な$O(NM)$畳み込みで十分。
2Garnerのアルゴリズムで復元
各係数について3つの余りから$p_1p_2p_3$を法とした一意な整数値を逐次的なmod逆元計算で復元。
各係数について3つの余りから$p_1p_2p_3$を法とした一意な整数値を逐次的なmod逆元計算で復元。
3最終的なMODでの還元
復元した値を最後に問題指定の$\mathrm{MOD}$で割った余りを取る。2段階を混同しないことが重要。
復元した値を最後に問題指定の$\mathrm{MOD}$で割った余りを取る。2段階を混同しないことが重要。
4計算量の見積もり
Step1をNTTに置き換えれば$O((N+M)\log(N+M))$まで高速化できる。
Step1をNTTに置き換えれば$O((N+M)\log(N+M))$まで高速化できる。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 3つの結果を平均や単純比較で合成しようとする | CRTの仕組みを理解せずアドホックに処理 | 必ずGarnerのアルゴリズムで復元する |
| Garner合成結果にさらにMOD還元するのを忘れる | 復元値が最終答えだと勘違い | 復元値は「真の畳み込み値」、最後に必ずMODで割った余りを出力 |
| 互いに素でない素数を選んでしまう | p1,p2,p3の互いに素性を確認せず選定 | 本問で指定された3素数(互いに素が保証済み)をそのまま使う |
| 真の畳み込み値が$p_1p_2p_3$を超える大きな入力を扱う | N,Mや係数の上限を考慮せず値域チェックを怠る | 制約範囲内なら問題ないが、大きなN,Mでは値域を確認する |
次のステップ
- 発展: Step1をNTTに置き換え、N,Mを$10^5$オーダーまで拡張して速度比較する
- 発展: 4つ以上の法を使う一般化されたCRT合成($k$個の法からの一般Garnerのアルゴリズム)を実装する
- 次回予告: Master Levelローテーション継続(次回テーマは実行時に選定)