部分和問題（N個の配列から和がKになるように選ぶ）とその解き方

2020年9月23日動的計画法bit全探索,DP,配列,部分和問題

部分和問題とは、\(N\) 個の数 \( a_1, a_2, …, a_n\) が与えられたとき、その中からいくつかを選んで和をちょうど \(K\) にできるか判定をする問題です。

例えば配列が [1,3,8] で、\(K\) が 9 のときは、1,8 を選ぶと和が 9 になり K に一致します。

全探索による単純なアルゴリズムでは計算量が大きくなってしまいますが、\(K\)と \(N\) の値が大きくなければ動的計画法で計算することができます。

アルゴリズム

※例えば配列が [1,3,8] で、\(K\) が 9 のときは以下のようにすべての選び方を探索します。

• 1, 3, 8 : 和は0
• 1, 3, 8 : 和は1
• 1, 3, 8 : 和は3
• 1, 3, 8 : 和は8
• 1, 3, 8 : 和は4
• 1, 3, 8 : 和は9
• 1, 3, 8 : 和は11
• 1, 3, 8 : 和は12

以上より和が 9 となり、\(K\) と一致するときがあると判定できます。

※動的計画法は、途中までの計算結果を利用することで余計な計算を省略する方法です。

※実際には前の計算結果の部分を上書きして計算することも可能なので、\(N \times K\) の二次元配列ではなく、要素数 \(K\) の1次元配列だけでも十分です。

計算量

• 全探索によるアルゴリズム
  • \(O(N 2^N)\)
• 動的計画法によるアルゴリズム
  • \(O(NK)\)

全探索をする場合は、\(2^N\) 通りについて、\(N\) 回足し算をするかチェックするので全体では \(O(N 2^N)\) になります。

動的計画法では、ループを用いて \(N \times K\) 回だけ配列上で計算を行うので \(O(NK)\) です。

ちなみに全探索をする場合は、再帰関数などを用いて工夫して足し算を行うことで、\(O(2^N)\) で計算することも可能です。

C++での実装例

以下のような入力が与えられたときのプログラムを考えます。

入力例:
N K
\(a_1 a_2 … a_N\)

全探索によるアルゴリズム

特に工夫をしない通常のビット全探索で計算ができます。

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

int main() {
    // 入力
    int N, K;
    cin >> N >> K;
    vector<int> a(N);
    for (int i = 0; i < N; i++) {
        cin >> a[i];
    }

    // ビット全探索
    for (int bits = 0; bits < (1 << N); bits++) {
        int sum = 0;
        for (int i = 0; i < N; i++) {
            if (bits & (1 << i)) {
                sum += a[i];
            }
        }
        if (sum == K) {
            cout << "Yes" << endl;
            return 0;
        }
    }

    cout << "No" << endl;
    return 0;
}

動的計画法

二次元配列を利用した動的計画法です。

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

int main() {
    // 入力
    int N, K;
    cin >> N >> K;
    vector<int> a(N);
    for (int i = 0; i < N; i++) {
        cin >> a[i];
    }

    // 二次元配列
    vector<vector<bool>> dp(N + 1, vector<bool>(K + 1, false));

    // 初期化
    for (int i = 0; i <= N; i++) {
        dp[i][0] = true;
    }

    // 更新
    for (int i = 0; i < N; i++) {
        for (int k = 0; k <= K; k++) {
            if (k - a[i] >= 0) dp[i + 1][k] = dp[i + 1][k] | dp[i][k - a[i]];
            dp[i + 1][k] = dp[i + 1][k] | dp[i][k];
        }
    }

    if (dp[N][K]) {
        cout << "Yes" << endl;
    } else {
        cout << "No" << endl;
    }
    return 0;
}