C++で合計が等しい2つの集合の最大合計を求めるプログラム

問題概要

数値のリスト nums が与えられます。このリストの要素を振り分けて合計が等しくなる2つの集合を作成し、そのときの合計値の最大を求めます。

たとえば、入力が nums = [2, 5, 4, 6] の場合、出力は 6 になります。[2, 4] と [6] という2つの集合を作れば、どちらも合計が 6 となり、これより大きい等しい合計は実現できないためです。

解決アプローチ(動的計画法)

各要素には「集合Aに入れる」「集合Bに入れる」「どちらにも入れない」という3つの選択肢があります。素朴な全探索では O(3^n) の計算量が必要になりますが、2つの集合の合計の差分を状態として管理する動的計画法(DP)を用いれば、効率的に解くことができます。

DPテーブル dp[i][j] は、「先頭から i 個の要素を処理したとき、2つの集合の合計差が j - sum となる状態における、片方の集合の合計の最大値」を表します。到達できない状態は -1 で表し、初期状態は dp[0][sum] = 0(差分ゼロ・合計ゼロ)とします。

アルゴリズムの手順

• sum := 0 とし、nums の各要素 i について sum := sum + i を実行する(総和を計算)
• n := nums のサイズ
• サイズ (n + 1) × (2 * sum + 5) の2次元配列 dp を定義し、すべて -1 で埋める
• dp[0][sum] := 0 とする
• i を 1 から n まで 1 ずつ増やしながら繰り返す:
  • x := nums[i - 1]
  • j を 0 から 2 * sum + 5 未満まで 1 ずつ増やしながら繰り返す:
    • j - x >= 0 かつ dp[i - 1][j - x] != -1 の場合:
      dp[i][j] := max(dp[i][j], dp[i - 1][j - x] + x)(要素 x を合計の大きい側の集合に追加)
    • j + x < 2 * sum + 5 かつ dp[i - 1][j + x] != -1 の場合:
      dp[i][j] := max(dp[i][j], dp[i - 1][j + x])(要素 x をもう一方の集合に追加)
    • dp[i][j] := max(dp[i][j], dp[i - 1][j])(要素 x を使用しない)
• dp[n][sum] を返す(合計差がゼロ、つまり2つの集合の合計が等しいときの最大値)

C++の実装例

理解を深めるために、以下の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int solve(vector<int>& nums) {
      int sum = 0;
      for (int i : nums) sum += i;
      int n = nums.size();
      vector<vector<int> > dp(n + 1, vector<int>(2 * sum + 5, -1));
      dp[0][sum] = 0;
      for (int i = 1; i <= n; i++) {
         int x = nums[i - 1];
         for (int j = 0; j < 2 * sum + 5; j++) {
            if (j - x >= 0 && dp[i - 1][j - x] != -1) {
               dp[i][j] = max(dp[i][j], dp[i - 1][j - x] + x);
            }
            if (j + x < 2 * sum + 5 && dp[i - 1][j + x] != -1) {
               dp[i][j] = max(dp[i][j], dp[i - 1][j + x]);
            }
            dp[i][j] = max(dp[i][j], dp[i - 1][j]);
         }
      }
      return dp[n][sum];
   }
};
int solve(vector<int>& nums) {
   return (new Solution())->solve(nums);
}
main(){
   vector<int> v = {2, 5, 4, 6};
   cout << solve(v);
}

入力

{2, 5, 4, 6}

出力

6

計算量の目安

• 時間計算量: O(n × sum) — 各要素ごとに差分の状態全体を走査します。
• 空間計算量: O(n × sum) — DPテーブルの保存に必要なメモリ量です。