Day 059-Q1 — XOR畳み込み (Walsh-Hadamard変換)

問題

$N$ 個の非負整数 $a_0, a_1, \dots, a_{N-1}$（$0 \le a_i < 2^M$）が与えられる。次の値を $\mathrm{mod}\ 998244353$ で求めよ。

ただし $\bigoplus$ はビット XOR、空集合の XOR は $0$ とする。

入出力例

3 2
1 2 3

30

$a = [1,2,3]$。部分集合 $\emptyset(0), \{1\}(1), \{2\}(2), \{3\}(3), \{1,2\}(3), \{1,3\}(2), \{2,3\}(1), \{1,2,3\}(0)$ の XOR の2乗和 = $0+1+4+9+9+4+1+2 = 30$。

概念図: Walsh-Hadamard 変換

ヒント（段階的開示）

def wht(a, inv=False):
    n = len(a)
    h = 1
    while h < n:
        for i in range(0, n, h*2):
            for j in range(h):
                x, y = a[i+j], a[i+j+h]
                a[i+j], a[i+j+h] = (x+y)%MOD, (x-y)%MOD
        h <<= 1
    if inv:
        inv_n = pow(n, MOD-2, MOD)
        for i in range(n): a[i] = a[i]*inv_n%MOD

f = [0]*SIZE; f[0] = 1
wht(f)
for a in A:
    delta = [0]*SIZE; delta[a] = 1
    wht(delta)
    for i in range(SIZE): f[i] = f[i]*(delta[i]+1)%MOD
wht(f, inv=True)
ans = sum(v*v*f[v] for v in range(SIZE)) % MOD

模範解答 (Python)

コピーimport sys
input = sys.stdin.readline

MOD = 998244353

def wht(a, inv=False):
    n = len(a)
    h = 1
    while h < n:
        for i in range(0, n, h*2):
            for j in range(h):
                x, y = a[i+j], a[i+j+h]
                a[i+j], a[i+j+h] = (x+y)%MOD, (x-y)%MOD
        h <<= 1
    if inv:
        inv_n = pow(n, MOD-2, MOD)
        for i in range(n):
            a[i] = a[i]*inv_n%MOD

def main():
    N, M = map(int, input().split())
    A = list(map(int, input().split()))
    SIZE = 1 << M

    f = [0]*SIZE
    f[0] = 1  # 空集合
    wht(f)

    for a in A:
        delta = [0]*SIZE
        delta[a] = 1
        wht(delta)
        for i in range(SIZE):
            f[i] = f[i]*(delta[i]+1)%MOD

    wht(f, inv=True)

    ans = 0
    for v in range(SIZE):
        ans = (ans + v*v%MOD*f[v]) % MOD
    print(ans)

main()

Step-by-Step 解説

Step 1: 問題の変換

$$\sum_S (\mathrm{XOR}_S)^2 = \sum_{v=0}^{2^M-1} v^2 \cdot f[v]$$ ここで $f[v] = |\{S : \mathrm{XOR}_S = v\}|$。$f$ を求めることが目標。

Step 2: XOR 畳み込みによる $f$ の更新

初期 $f = \mathbf{1}_0$（空集合のみ）。各要素 $a_i$ を加えると： $f^{\text{new}}[v] = f[v] + f[v \oplus a_i]$（$a_i$ を選ぶ/選ばない）。 これを全要素に適用 = $f$ と $\mathbf{1}_{a_i}$ の XOR 畳み込みの "加算版"。

Step 3: WHT の活用

WHT ドメインでは XOR 畳み込みが点ごとの積に変換される。$\hat{f} \cdot (1 + \hat{\delta}_{a_i})$ を全要素に適用後、逆 WHT で $f$ を復元。

Step 4: 計算量

WHT: $O(2^M \cdot M)$。全要素の delta 適用: $O(N \cdot 2^M)$。合計 $O(N \cdot 2^M + 2^M \cdot M)$。$N = 2 \times 10^5, M = 20$ で約 $4 \times 10^7$ 演算。

よくあるミス

次のステップ

発展問題: $\sum_S (\mathrm{XOR}_S = T)$ を満たす部分集合数を全 $T$ について求めよ（XOR Basis 応用との比較）。

自己評価