C++で配列をK個の等しい合計値の部分集合に分割できるか判定する方法

問題概要

整数型の配列 nums と正整数 k が与えられます。この配列を、各部分集合の合計値がすべて等しくなるような k 個の空でない部分集合に分割できるかどうかを判定してください。

例えば、配列が [4,3,2,3,5,2,1]、k = 4 の場合を見てみましょう。この配列は [[5], [1,4], [2,3], [2,3]] のように4つの部分集合に分割でき、それぞれの合計はすべて 5 で等しくなります。したがって、答えは True となります。

解法のポイント：ビットマスクDP

この問題は、ビットマスク動的計画法（bitmask DP）によって効率的に解くことができます。n 個の要素に対して 2^n 通りの「使用済み要素の組み合わせ」を状態として持ち、各状態における累積合計を記録しながら探索を進めます。また、配列をあらかじめソートしておくことで、現在構築中のサブセットに入らない要素が出た時点で探索を打ち切れる点も重要な工夫です。

アルゴリズムの手順

1. サイズ 2^n の2つのテーブル dp（到達可能フラグ）と total（累積合計）を用意します。
2. 配列 nums を昇順にソートし、全要素の合計を sum とします。
3. sum % k が 0 以外、またはソート後の最大要素が sum / k を超える場合は false を返します（この時点で分割不可能と確定するため）。
4. dp[0] = true とし、sum = sum / k（＝各サブセットの目標合計値）を設定します。
5. i を 0 から 2^n − 1 まで走査します。dp[i] が true の状態について、未使用の要素 j を追加できるかを試します。
   • 新しい状態 temp = i | (1 << j) を作成します。
   • nums[j] ≤ sum − (total[i] % sum) を満たす場合は要素を追加でき、dp[temp] = true、total[temp] = total[i] + nums[j] と更新します。
   • 条件を満たさない場合は、配列がソート済みであることからそれ以降のより大きな要素も収まらないため、内側のループを break します。
6. 最終的に dp[(2^n) − 1]（全要素を使用した状態）が true であれば true を返します。

ここで total[i] % sum は「現在構築途中のサブセットの合計」を表します。次の要素を加えても目標値 sum を超えない場合のみ、その遷移が有効になります。

C++による実装例

以下の実装を見ると、仕組みがより理解しやすくなります。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   bool canPartitionKSubsets(vector<int>& nums, int k) {
      int n = nums.size();
      vector <bool> dp(1 << n);
      vector <int> total(1 << n);
      sort(nums.begin(), nums.end());
      int sum = 0;
      for(int i = 0; i < nums.size(); i++)sum += nums[i];
      if(sum % k || nums[nums.size() - 1] > sum / k) return false;
      dp[0] = true;
      sum /= k;
      for(int i = 0; i < (1 << n); i++){
         if(dp[i]){
            for(int j = 0; j < n; j++){
               int temp = i | (1 << j);
               if(temp != i){
                  if(nums[j] <= sum - (total[i] % sum)){
                     dp[temp] = true;
                     total[temp] = total[i] + nums[j];
                  }
                  else{
                     break;
                  }
               }
            }
         }
      }
      return dp[(1 << n) - 1];
   }
};
main(){
   Solution ob;
   vector<int> v = {4,3,2,3,5,2,1};
   cout << (ob.canPartitionKSubsets(v, 4));
}

入力

[4,3,2,3,5,2,1]
4

出力

1

出力が 1（true）となり、この配列が合計値の等しい4つの部分集合に分割可能であることが確認できます。

計算量

• 時間計算量： O(2^n × n) — 状態数 2^n に対して、各状態から最大 n 個の遷移を調べます。
• 空間計算量： O(2^n) — dp テーブルと total テーブルの保存に必要です。