高速ゼータ変換

これは データ構造とアルゴリズム 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$ になっています。

前と同じように愚直に解こうとすると $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) & ii になってしまうので 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(2^NN)$ で解くことができます。

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

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

まず、dp[i][0] = $A_i$ になります。

そして、

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

・$i$ の下から $j$ 番目 (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つの整数を持つ場合です)

日本情報オリンピック 2018年本選 - 5 毒蛇の脱走 (高速ゼータ変換の他に包除原理というテクニックを使います、とても難しい)