FWTを用いたビット演算畳み込み

概要

高速ウォルシュ変換(FWT)は、FFTが通常の畳み込みを効率的に計算するのと同様に、ビット演算に基づく畳み込み演算を高速化するための手法です。ビット演算畳み込みでは、二つの数列の要素どうしがビット演算の結果に対して寄与します。

具体的には、数列AとBに対して、S[k] = Σi⊕j=k A[i]×B[j] を計算します。ここで⊕はOR、AND、XORなどのビット演算を表します。

変換の基本原理

FWTは線形変換であり、変換係数c(i,j)を導入します。FWT(A)[i] = Σj=0n-1 c(i,j)×A[j] と表現できます。

FWT(A)・FWT(B) = FWT(C) が成り立つためには、c(i,j)×c(i,k) = c(i,j⊕k) という条件を満たす必要があります。ビット演算の独立性により、各ビットを個別に考慮でき、2×2の行列でこの条件を実現できます。

基本変換行列

代表的な三つの演算に対する変換行列は以下の通りです:

  • OR畳み込み
    正変換: [[1,0],[1,1]]
    逆変換: [[1,0],[-1,1]]
  • AND畳み込み
    正変換: [[1,1],[0,1]]
    逆変換: [[1,-1],[0,1]]
  • XOR畳み込み
    正変換: [[1,1],[1,-1]]
    逆変換: [[1/2,1/2],[1/2,-1/2]]

FWTアルゴリズム

配列を上位ビットで分割する分割統治法を採用します。長さnの配列を前半と後半に分け、それぞれに対して再帰的にFWTを適用した後、2×2行列を用いて結果を合成します。

合成の際、前半の要素と後半の要素に対して行列の係数を乗じて加算することで、全体の変換結果を得ます。計算量はO(n log n)です。

実装例

#include <bits/stdc++.h>
using namespace std;

const int MOD = 998244353;
const int INV2 = (MOD + 1) / 2;
const int MAX_BITS = 20;

using ll = long long;

void walsh_transform(vector<ll>& vec, const ll mtx[2][2]) {
    int n = vec.size();
    for (int step = 1; step < n; step <<= 1) {
        for (int i = 0; i < n; i += (step << 1)) {
            for (int j = 0; j < step; ++j) {
                ll u = vec[i + j];
                ll v = vec[i + j + step];
                vec[i + j] = (mtx[0][0] * u + mtx[0][1] * v) % MOD;
                vec[i + j + step] = (mtx[1][0] * u + mtx[1][1] * v) % MOD;
            }
        }
    }
}

void apply_convolution(vector<ll> a, vector<ll> b, 
                       const ll fwd[2][2], const ll inv[2][2]) {
    int n = a.size();
    walsh_transform(a, fwd);
    walsh_transform(b, fwd);
    for (int i = 0; i < n; ++i) a[i] = a[i] * b[i] % MOD;
    walsh_transform(a, inv);
    for (int i = 0; i < n; ++i) printf("%lld ", a[i]);
    printf("\n");
}

int main() {
    int m; scanf("%d", &m);
    int n = 1 << m;
    
    vector<ll> input_a(n), input_b(n);
    for (ll& x : input_a) scanf("%lld", &x);
    for (ll& x : input_b) scanf("%lld", &x);
    
    const ll OR_F[2][2] = {{1, 0}, {1, 1}};
    const ll OR_I[2][2] = {{1, 0}, {MOD - 1, 1}};
    
    const ll AND_F[2][2] = {{1, 1}, {0, 1}};
    const ll AND_I[2][2] = {{1, MOD - 1}, {0, 1}};
    
    const ll XOR_F[2][2] = {{1, 1}, {1, MOD - 1}};
    const ll XOR_I[2][2] = {{INV2, INV2}, {INV2, MOD - INV2}};
    
    apply_convolution(input_a, input_b, OR_F, OR_I);
    apply_convolution(input_a, input_b, AND_F, AND_I);
    apply_convolution(input_a, input_b, XOR_F, XOR_I);
    
    return 0;
}

部分集合畳み込み

部分集合畳み込みは、i & j = 0 かつ i | j = k を満たす組み合わせに対して計算を行います。ポピュレーションカウント(1のビット数)を利用することで効率的に処理できます。

各ビット数ごとに別々の配列を用意し、それぞれに対してOR畳み込みを行った後、ビット数の次元で畳み込みを合成します。計算量はO(m²2^m)です。

部分集合畳み込みの実装

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1000000009;
const int MAX_LOG = 20;
const int MAX_N = 1 << MAX_LOG;

int bit_count(int x) { return __builtin_popcount(x); }

void fwt_subset(int arr[], const int mat[2][2], int len) {
    for (int seg = 2; seg <= len; seg <<= 1) {
        int half = seg >> 1;
        for (int start = 0; start < len; start += seg) {
            for (int offset = 0; offset < half; ++offset) {
                int idx1 = start + offset;
                int idx2 = start + offset + half;
                int val1 = arr[idx1];
                int val2 = arr[idx2];
                arr[idx1] = (mat[0][0] * val1 + mat[0][1] * val2) % MOD;
                arr[idx2] = (1LL * mat[1][0] * val1 + mat[1][1] * val2) % MOD;
            }
        }
    }
}

int main() {
    int m; scanf("%d", &m);
    int total = 1 << m;
    
    static int layerA[MAX_LOG + 1][MAX_N];
    static int layerB[MAX_LOG + 1][MAX_N];
    static int layerC[MAX_LOG + 1][MAX_N];
    
    for (int i = 0; i < total; ++i) {
        scanf("%d", &layerA[bit_count(i)][i]);
    }
    for (int i = 0; i < total; ++i) {
        scanf("%d", &layerB[bit_count(i)][i]);
    }
    
    const int FWD_MAT[2][2] = {{1, 0}, {1, 1}};
    const int INV_MAT[2][2] = {{1, 0}, {MOD - 1, 1}};
    
    for (int k = 0; k <= m; ++k) {
        fwt_subset(layerA[k], FWD_MAT, total);
        fwt_subset(layerB[k], FWD_MAT, total);
    }
    
    for (int k = 0; k <= m; ++k) {
        for (int i = 0; i <= k; ++i) {
            for (int pos = 0; pos < total; ++pos) {
                layerC[k][pos] = (layerC[k][pos] + 1LL * layerA[i][pos] * layerB[k - i][pos]) % MOD;
            }
        }
    }
    
    for (int k = 0; k <= m; ++k) {
        fwt_subset(layerC[k], INV_MAT, total);
    }
    
    for (int i = 0; i < total; ++i) {
        printf("%d ", layerC[bit_count(i)][i]);
    }
    printf("\n");
    
    return 0;
}

FWTによる高速累乗計算

FWT変換後の各要素に対して累乗計算を行い、逆変換することで、多項式のt乗を効率的に求められます。これは畳み込み演算の繰り返しを避ける手法として有効です。

タグ: FWT 畳み込み ビット演算 高速変換 部分集合畳み込み

8月1日 09:20 投稿