問題
長さ $2^n$ の数列 $A_0,\dots,A_{2^n-1}$ と $B_0,\dots,B_{2^n-1}$ が与えられる($0\le n\le18$)。
以下で定義される数列 $C_0,\dots,C_{2^n-1}$ を $998244353$ で割った余りとして求めよ。
$$C_k = \sum_{i \oplus j = k} A_i B_j \pmod{998244353}$$
$\oplus$はビット単位XOR。総和は$0\le i,j<2^n$かつ$i\oplus j=k$を満たす全ての組$(i,j)$について取る。
Walsh-Hadamard変換(WHT)を用いること。長さ$2^n$の配列$a$に対する(正規化前の)WHTは、$h=1,2,4,\dots,2^{n-1}$の順に、各ブロック$[i,i+h)$と$[i+h,i+2h)$のペア$(a_j,a_{j+h})$を$(a_j+a_{j+h},\ a_j-a_{j+h})$に同時更新するバタフライ演算を繰り返すことで計算できる。$A,B$それぞれにWHTを適用したものを要素ごとに掛け合わせ、逆WHT(WHTと同じ操作の後、全要素を$2^n$で割る)を適用すると$C$が得られる。
入力形式
n
A_0 A_1 ... A_{2^n-1}
B_0 B_1 ... B_{2^n-1}
制約
入出力例
入力例1
2
3 1 4 1
5 9 2 6出力例1
38 58 41 61入力例2
3
1 2 3 4 5 6 7 8
8 7 6 5 4 3 2 1出力例2
120 124 136 140 184 188 200 204概念図
ヒント(段階的開示)
ヒント1(方向性)
AND畳み込み・OR畳み込みが部分集合ゼータ/メビウス変換という「単調な包含関係」に基づく変換を使うのに対し、XOR畳み込みは全く異なる直交変換であるWalsh-Hadamard変換を使う。WHTは「$\pm1$の$2\times2$直交行列をテンソル積的に$n$回重ねたもの」を配列に掛ける操作に相当し、変換後の空間ではXOR畳み込みが単なる要素ごとの積になる(畳み込み定理のXOR版)。
ヒント2(アプローチ)
FFTやOR畳み込みのゼータ変換と同様に、WHTも「サイズ2のブロックごとの変換をブロック幅を倍々にしながら繰り返す」分割統治的な構造を持ち$O(2^n\cdot n)$で計算できる。バタフライ演算$(a_j,a_{j+h})\to(a_j+a_{j+h},a_j-a_{j+h})$を全ブロック・全レベルに適用するだけで変換が完了する。逆変換は同じ演算をもう一度行い、最後に配列全体を$2^n$で割ればよい。
ヒント3(誘導)
MOD = 998244353
def wht(a, invert):
n = len(a)
h = 1
while h < n:
for i in range(0, n, h * 2):
for j in range(i, i + h):
x, y = a[j], a[j + h]
a[j] = (x + y) % MOD
a[j + h] = (x - y) % MOD
h *= 2
if invert:
inv_n = pow(n, MOD - 2, MOD)
for i in range(n):
a[i] = a[i] * inv_n % MOD
def xor_convolution(A, B):
a, b = A[:], B[:]
wht(a, False); wht(b, False)
c = [a[i] * b[i] % MOD for i in range(len(a))]
wht(c, True)
return c
invert=Trueのときの割り算はmod逆元(フェルマーの小定理)で行う。
模範解答 (Python)
import sys
MOD = 998244353
def wht(a, invert):
n = len(a)
h = 1
while h < n:
for i in range(0, n, h * 2):
for j in range(i, i + h):
x = a[j]
y = a[j + h]
a[j] = (x + y) % MOD
a[j + h] = (x - y) % MOD
h *= 2
if invert:
inv_n = pow(n, MOD - 2, MOD)
for i in range(n):
a[i] = a[i] * inv_n % MOD
def xor_convolution(A, B):
n = len(A)
a = A[:]
b = B[:]
wht(a, False)
wht(b, False)
c = [a[i] * b[i] % MOD for i in range(n)]
wht(c, True)
return c
def main():
data = sys.stdin.read().split()
ptr = 0
n = int(data[ptr]); ptr += 1
size = 1 << n
A = [int(data[ptr + i]) % MOD for i in range(size)]; ptr += size
B = [int(data[ptr + i]) % MOD for i in range(size)]; ptr += size
C = xor_convolution(A, B)
print(' '.join(map(str, C)))
main()
Step-by-Step 解説
$\hat{A}_k=\sum_i(-1)^{\mathrm{popcount}(i\&k)}A_i$と定義すると、$(-1)^{\mathrm{popcount}(i\&k)}\cdot(-1)^{\mathrm{popcount}(j\&k)}=(-1)^{\mathrm{popcount}((i\oplus j)\&k)}$というXORとpopcountパリティの相性の良さから畳み込み定理$\widehat{A*B}=\hat{A}\cdot\hat{B}$が成立する。
$n$変数の$\pm1$直交変換はテンソル積$H^{\otimes n}$で表現でき、愚直な行列積は$O(4^n)$だが、FFTと同様に1変数ずつ順番に適用することで$O(2^n\cdot n)$に削減できる。
$H^2=2I$という性質から、順変換と全く同じバタフライ演算をもう一度適用し全要素を$2^n$で割れば逆変換になる。mod演算下ではフェルマーの小定理による逆元計算を使う。
OR畳み込み(Day116)は部分集合ゼータ変換、AND畳み込み(Day114)はその双対、XOR畳み込みはWalsh-Hadamard変換を使う。3つとも「$n$段のビット単位バタフライ演算」という共通の骨格を持つ。
WHTは$O(2^n\cdot n)$、全体でも順変換2回・逆変換1回・要素積1回で$O(2^n\cdot n)$。$n\le18$でも約470万回程度のバタフライ演算で十分高速。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| バタフライ演算の減算結果が負になりmodを忘れる | Pythonの%は自動補正されるが他言語では明示補正が必要なことを見落とす | Pythonでは問題ないが、C++等に移植する際は(x-y+MOD)%MODにする |
| 逆変換で割る数を固定値2と勘違いする | 「1回のバタフライ演算につき2で割る」という誤解 | 逆変換の正規化は最後にまとめて$2^n$で1回割る |
| OR/AND畳み込み用のゼータ変換コードを流用し加算のみの更新にしてしまう | 3つの畳み込みの変換内容の違いを混同 | XOR畳み込みは必ず「和と差」のバタフライ演算を使う |
| $n=0$(配列長1)のケースを特別扱いせず境界エラーになる | ループ条件h<nが最初から不成立のケースを想定していない | ループが1回も回らず恒等変換になるのが正しい挙動で特別扱い不要 |
| 入力の$A_i,B_i$をmod済みとして扱わず不要な多倍長演算コストをかける | 読み込み直後のmod処理を省略 | 読み込み時に% MODを必ず適用する |
次のステップ
- 発展: 本問のXOR畳み込みを使い、$N$個の非負整数の部分集合のうち「XORが$0$になる部分集合の個数」を数える問題(線形基底を使う別解と結果が一致することを確認せよ)に応用せよ。
- 発展2: Day114(AND畳み込み)・Day116(OR畳み込み)・本問(XOR畳み込み)の3つの変換を1つのライブラリにまとめ、切り替え可能なインターフェースを実装せよ。