部分集合内を1回ずつ見るビットdpを高速に計算する

まえがき

CSA Round #74 5問目 Good Permutation を解きました。コンテスト中には考察はほぼできていたのですが、ある部分での高速化ができず通すことができませんでした。その部分についての知見を頂いたので、記事にして残しておこうと思います。 なんかそれっぽい部分は2進数だと思って読んでください。

似たようで簡単な問題

次のようなビットDPを計算することを考えます。ビットを集合としてみる場合、ビットが立っているとき対応する要素が集合に含まれているものとみなします。またsub(i)をiを集合としてみたときの部分集合の集合とします。(ex. sub(101)={101,100,001,000})
ビット長はnとします。
A:長さ2^nの配列が与えられて、
dp[i]=max{k∈sub(i)}(A[k]))
これはi=0から初めて、iの立っている各ビットを1つずつ0にしたもの全てとmaxをとる遷移をすればよいです。以下のようなコードで実現できます。

for (int i = 0; i < (1LL << n); ++i)dp[i] = A[i];

for (int i = 0; i < (1 << n); ++i) {
  for (int j = 0; j < n; ++j) {
		if (i&(1 << j)) {//iのjビット目が立っているなら
			dp[i] += dp[i ^ (1 << j)];
		}
	}
}

本題

次のようなビットDPを計算します。
A:長さの2^nの配列  
dp[i]=\sum_{k∈sub(i)}^{}A[k]
前の節でのdpの定義のmaxの部分がsumに変わっただけですね。しかしこれは先程と同じようにして解くことはできません。具体例を用いて理由を説明します。 先程の問題でdpを計算する場合、dp[101]=dp[100]+dp[001]+A[101]=dp[000]+A[100]+dp[000]+A[001]+A[101]=2*A[000]+A[001]+A[100]+A[101]となり、A[000]が2回足されてしまいます。iに立っているビットが複数ある場合に、どのビットから0にしていくかの順序が変わると別の遷移ができてしまい、複数回足されることになってしまいます。(例えば 101→100→000と101→001→000) 前の問題では演算がmaxで、max(a,a)=aであったために問題はなかったのですがsum(a,a)!=aではないため、各A[i]が複数回足されると困ってしまいます。
まず愚直な解法として、毎回部分集合を列挙して足す方法を考えてみます。これは以下のようなコードになります。
(配列名をdpのままにしているけれどこの解法はdpではないです)

for (int i = 0; i < (1<<n); ++i) {
	for (int j = i;; j = (j - 1)&i) {
		dp[i] += a[j];
		if (j <= 0)break;
	}
}

内側のループは、iの部分集合を列挙しています。iの部分集合の数は2^{popcnt(i)}個あるので、内側のループはpopcnt(i)回回ります。
従って全体でのループ回数は \sum_{i=0}^{2^n} 2^{popcnt(i)}です。
popcnt(i)==kとなるiは_n C _k個あるので、これは\sum_{k=0}^{n} {_n C _k 2^k}と等しいです。
さらにこれは (1+2)^nを展開した形になっているので、全体のループ回数は3^nであることが分かります。
少し遅いですね。n=20が通らないぐらいでしょうか。

実はこのdpは O(n*2^n)で計算することができます。
結論から述べますが、これは以下のようなc++のコードで実現できます。

for (int i = 0; i < (1LL << n); ++i)dp[i] = A[i];

for (int i = 0; i < n; ++i) {
	for (int j = 0; j < (1 << n); ++j) {
		if (j&(1 << i)) {//jのiビット目が立っているなら
			dp[j] += dp[j ^ (1 << i)];
		}
	}
}

コードを見ると、iとjのループの順番が入れ替わっているだけです。これでうまくいくことの証明を帰納法で行います。

証明

f(j,i)=jの右からiビット(1-indexed)はjの部分集合になっていてかつjの左からn-iビットはjと等しいような整数全体の集合
とします。 (ex.f(1011,2)={1011,1010,1001,1000})
示すべき性質は、
「外側のループがi回終わったとき、dp[j]の値は、全てのk∈f(j,i)についてA[k]の和をとったものと等しい」です。この性質をPと呼びます。
i=nのときPが成り立てば、このアルゴリズムの正当性が示せます。

i=0の場合

f(j,0)={j}であり、ループに入る前はdp[j]=A[j]であるため、i=0のときはPは成立します。

i=mのとき成立することを仮定し、i=m+1のとき

(1)jのm+1ビット目が0のとき
まず、f(j,m)=f(j,m+1)が成り立つことがわかります。
さらにjのm+1ビット目が0である場合、m+1回目のループではdp[j]の値は変わらないことがわかります。従って仮定よりPが満たされることが示せます

(2)jのm+1ビット目が1の場合
f(j,m+1)は、f(j,m)全体と、∀ k∈f(j,m)のm+1ビット目を0に
したもの全体の集合の合併となるから、
f[j,m+1)=f(j,m)∪f(j^2^m,m)が成り立ちます。ここで^は排他的論理和です。
さらに、f(j,m)とf(j^2^m,m)は排反です。これは、f(j,m)の任意の要素のm+1ビット目は1であり、f(j^2^m,m)の任意の要素のm+1ビット目は0であることから分かります。
一方、m+1回目のループではdp[j]の値にdp[j^2^m]のみが加算されています。j^2^mのm+1ビット目は0であることを考えると、dp[j^2^m]は∀ k∈f(j^2^m,m)についてのA[k]の和と等しく、dp[j]は∀k∈f(j,m)についてのA[k]の和と等しいです。
したがって、m+1回目のループ終了時ではdp[j]は∀k∈f(j,m+1)についてのA[k]の和と等しいので、i=m+1のときもPが成立することが分かります。

以上から、数学的帰納法より∀i 1<=i<=nについてPが成り立つことが分かります。
つまりこのアルゴリズムの正当性も示せました。

おわりに

競プロ的には通るnの範囲がn<=18ぐらいからn<=25ぐらいに広がるぐらいの差ですが、実際にこの差で通らなかったので意義はあると思い記事にしました。実はもっと早くなったり、間違いがあった場合は教えてくださるとうれしーです。