サブセットの総和 – C++での動的計画法による効率的な解法

問題概要

この問題では、サイズ 2ⁿ の配列 arr[] が与えられます。目的は、動的計画法(DP)を用いて「サブセットの総和(Sum over Subsets)」を求めるプログラムを作成することです。

具体的には、次の関数 F(x) を計算します。

F(x) = Σ A_i （ただし x & i == i を満たすすべての i に対する和）

つまり、i が x のビット単位での部分集合であるような A_i の合計を求めるという意味です。

入力例と出力例

入力: A[] = {5, 7, 1, 9}, n = 2

出力: 5 12 6 22

説明: n = 2 のとき、x は 0、1、2、3 の 4 通りの値を取ります。それぞれの関数値は以下のように計算できます。

F(0) = A₀ = 5
F(1) = A₀ + A₁ = 5 + 7 = 12
F(2) = A₀ + A₂ = 5 + 1 = 6
F(3) = A₀ + A₁ + A₂ + A₃ = 5 + 7 + 1 + 9 = 22

解法のアプローチ

この問題を動的計画法で解くには、各マスク(mask)に着目し、それぞれのマスクについてビット単位の部分集合を効率よく求めます。部分集合の和を DP テーブルに記録しておくことで、重複する計算を大幅に削減できます。素朴な全探索では、ビットが立っているかどうかにかかわらず、各インデックスが 2ⁿ 個のマスクから何度も参照されてしまいます。

i 番目のビットの状態に応じて、次のような漸化式を立てます。

i 番目のビットが立っている場合：
DP(mask, i) = DP(mask, i-1) + DP(mask ^ 2ⁱ, i-1)

i 番目のビットが立っていない場合：
DP(mask, i) = DP(mask, i-1)

この手法は「SOS DP(Sum over Subsets DP)」として知られており、時間計算量は O(n ・ 2ⁿ)、空間計算量も O(n ・ 2ⁿ) となります。

ソリューションの実装例(C++)

#include <iostream>
using namespace std;
const int N = 1000;

void SumOverSubsets(int a[], int n) {

    int sum[1 << n] = {0};
    int DP[N][N];
    for (int i = 0; i < (1 << n); i++) {
        for (int j = 0; j < n; j++) {
            if (i & (1 << j)) {
                if (j == 0)
                    DP[i][j] = a[i] + a[i ^ (1 << j)];
                else
                    DP[i][j] = DP[i][j - 1] + DP[i ^ (1 << j)][j - 1];
            } else {
                if (j == 0)
                    DP[i][j] = a[i];
                else
                    DP[i][j] = DP[i][j - 1];
            }
        }
        sum[i] = DP[i][n - 1];
    }
    for (int i = 0; i < (1 << n); i++)
        cout<<sum[i]<<"\t";
}
int main() {
    int A[] = {5, 7, 1, 9};
    int n = 2;
    cout<<"The sum over subsets is \t";
    SumOverSubsets(A, n);
    return 0;
}

出力

The sum over subsets is 5 12 6 22