高速ゼータ変換について

これは データ構造とアルゴリズム Advent Calendar 2018 の 25 日目の記事です。

突然ですが

次の問題を解いてみましょう。

整数 $N$ と $A_0, A_1, A_2, ... A_{2^N-1}$ が与えられます。
$s, t$ に対して $s\ AND\ t=t$ の時、$t \subset s$ と表記することにします。
$0≦i<2^N$ を満たす整数 $i$ について、$j \subset i$ を満たす整数 $j$ についての $A_j$ の総和をそれぞれ求めなさい。
ただし、$a\ AND\ b$ は $a$ と $b$ のbitごとの論理積を求める演算とします。

$1≦N≦10$
$1≦A_i≦10^9\ (0≦i<2^N)$
実行時間制限 2秒

問題文にある $s\ AND\ t=t$ が何を言っているのかちょっと分かりにくいかもしれませんが、これは $t$ で立っている (=1になっている) bitは $s$ でも立っているということです。
例を示します。

$N=2, A=[1, 10, 100, 1000]$ のとき
$0≦i<2^N$ を満たす整数 $i$ は $0, 1, 2, 3$
$i=0$ のとき、$i\ AND\ j=j$ を満たす $j$ は $0$ しかないので、答えは $1$
$i=1$ のとき、$i\ AND\ j=j$ を満たす $j$ は $0,1$ なので、答えは $11$
$i=2$ のとき、$i\ AND\ j=j$ を満たす $j$ は $0,2$ なので、答えは $101$
$i=3$ のとき、$i\ AND\ j=j$ を満たす $j$ は $0,1,2,3$ なので、答えは $1111$

この問題は条件を愚直に試していくことで解くことができます。
$i$ の探索に $O(2^N)$ で、$j$ の探索に $O(2^N)$ の計算量が必要なので、全体の計算量はその積の $O(4^N)$ になります。
$N≦10$ なので $4^N≦1048576$ となり、実行時間制限の2秒には十分に間に合います。

for(int i = 0; i < (1 << N); i++){
    int ans = 0;
    for(int j = 0; j < (1 << N); j++){
        if((i & j) == j){
            ans += A[j];
        }
    }
    cout << ans << endl;
}

ちょっとしたテクニック

※この節は高速ゼータ変換とは関係がないので、この節を飛ばして次の"本題"の節を読んでもらっても大丈夫です。

では、次の問題はどうでしょうか。

整数 $N$ と $A_0, A_1, A_2, ... A_{2^N-1}$ が与えられます。
$0≦i<2^N$ を満たす整数 $i$ について、$j \subset i$ を満たす整数 $j$ についての $A_j$ の総和をそれぞれ求めなさい。

$1≦N≦15$
$1≦A_i≦10^9\ (0≦i<2^N)$
実行時間制限 2秒

前の問題と同じように見えますが、よく見ると制約が少し変わっています。
前の問題では $1≦N≦10$ だったのが $1≦N≦15$ になっています。
前と同じように愚直に解こうとすると $N=15$ のとき $4^N=1073741824$ となり、2秒の実行時間制限に間に合いません。

ここで、テクニックを使うことで高速に解くことができます。
先ほど

for(int j = 0; j < (1 << N); j++){
    if((i & j) == j){
        ans += A[j];
    }
}

としていた部分を

for(int j = i; j >= 0; j--){
    if((i & j) == j){
        ans += A[j];
    }
}

とします。int j = i の部分が気になるかもしれませんが、$j \subset i$ の時は必ず $j≦i$ なので変数 j の初期値を i にしても大丈夫です。
まだこれでは計算量は変わっていません。そして、これを

for(int j = i; j > 0; j = (j - 1) & i){
    ans += A[j];
}
ans += A[0];

とします。$j - 1$ は2進法で考えると、$j$ で立っているbitのうち最下位のものを消し、それ未満のbitを全て立てたものになります。
$j$ を自分で実際に決めて $j$ と $j - 1$ を2進法で表してみるとこれが分かると思います。

そして、(j - 1) & i は $j - 1$ の立っているbitのうち $i$ でも立っているものだけを取り出したものになります。
これも実際にfor文の動きを辿ってみると分かると思います。

ただし、j が 0 の場合、(0 - 1) は -1 になりますが、これはC++では内部的に 111..1 であり、(j - 1) & i は i になってしまうので j が 0 の場合はfor文の外で判定をしています。
また、上のコードでは $j \subset i$ を満たす $j$ をもれなくだぶりなく列挙しているので、2個上のコードにあったif文が必要無くなります。

このテクニックは計算量が少し分かりにくいですが、整数 $x$ の立っているbitの個数を $bitcount(x)$ と表記すると、上のコードのfor文のループ回数は $O(2^{bitcount(i)})$ になるので、問題全体の計算量は $O(3^N)$ になります。(二項定理で $(2+1)^N$ を展開すると分かる)
$N≦15$ なので $3^N≦14348907$ となり、実行時間制限に間に合います。

本題

では、今度は次の問題を解いてみましょう。

整数 $N$ と $A_0, A_1, A_2, ... A_{2^N-1}$ が与えられます。
$0≦i<2^N$ を満たす整数 $i$ について、$j \subset i$ を満たす整数 $j$ についての $A_j$ の総和をそれぞれ求めなさい。

$1≦N≦20$
$1≦A_i≦10^9\ (0≦i<2^N)$
実行時間制限 2秒

$N≦20$ なので $3^N$ は最大で $3486784401$ になり、$O(3^N)$ 時間のアルゴリズムでも間に合いません。
やっと本題になりますが、ここで 高速ゼータ変換 というアルゴリズムを使うことで $O(2^NN)$ 時間で解くことができます。

具体的には、次のようなDPをします。

• dp[i][j] $= k \subset i$ を満たし、かつ $k$ の下から $j$ 番目以上 (0-indexed) のbitの状態が $i$ のそれと同じような整数 $k$ についての $A_k$ の総和

まず、dp[i][0] $= A_i\ (0≦i<2^N)$ になります。
そして、

• $i$ の下から $j - 1$ 番目 (0-indexed) のbitが立っているなら dp[i][j] = dp[i & ~(1 << (j - 1))][j - 1] + dp[i][j - 1] $(1≦j≦N)$
• $i$ の下から $j - 1$ 番目 (0-indexed) のbitが立っていないなら dp[i][j] = dp[i][j - 1] $(1≦j≦N)$

と更新できます。
また、求めたい答えは dp[i][N] になっています。
よって、全体のコードはこうなります。

int dp[1 << N][N + 1];
for(int i = 0; i < (1 << N); i++){
    dp[i][0] = A[i];
}
for(int j = 0; j < N; j++){
    for(int i = 0; i < (1 << N); i++){
        if(i & (1 << j)){
            dp[i][j + 1] = dp[i & ~(1 << j)][j] + dp[i][j];
        }else{
            dp[i][j + 1] = dp[i][j];
        }
    }
}
for(int i = 0; i < (1 << N); i++){
    cout << dp[i][N] << endl;
}

次のように、DPテーブルを再利用することで使うメモリを $O(2^N)$ まで減らすことができます。

int dp[1 << N];
for(int i = 0; i < (1 << N); i++){
    dp[i] = A[i];
}
for(int j = 0; j < N; j++){
    for(int i = 0; i < (1 << N); i++){
        if(i & (1 << j)){
            dp[i] += dp[i & ~(1 << j)];
        }
    }
}
for(int i = 0; i < (1 << N); i++){
    cout << dp[i] << endl;
}

DPテーブルに配列 A を使うことで、初期化をしなくてもよくなります。

for(int j = 0; j < N; j++){
    for(int i = 0; i < (1 << N); i++){
        if(i & (1 << j)){
            A[i] += A[i & ~(1 << j)];
        }
    }
}
for(int i = 0; i < (1 << N); i++){
    cout << A[i] << endl;
}

計算量は $O(2^NN)$ になります。
$N≦20$ の時 $2^NN≦20971520$ なので、実行時間制限の2秒に間に合います。

終わりに

今回は総和でしたが、総和でなくても解ける場合があります。(交換法則と結合法則が成り立てば良い)
このアルゴリズムは競プロでは稀にしか出てきませんが、稀に出てくるので覚えておいて損はありません。

最後にこのアルゴリズムを使う問題を挙げておきます。

ネタバレ注意

AtCoder Regular Contest 100 - E Or Plus Max (DPテーブルの要素がそれぞれ2つの整数を持つ場合です)
No.767 配られたジャパリまん - yukicoder (高速ゼータ変換の他に包除原理というテクニックを使います)
日本情報オリンピック 2018年本選 - 5 毒蛇の脱走 (こちらも包除原理を使います、とても難しい)