C++でリストを合計が等しいk個の部分集合に分割できるか判定するプログラム

問題の概要

数値のリスト nums と整数 k が与えられたとき、nums を「各部分集合の合計がすべて等しくなる」ような k 個の部分集合に分割できるかどうかを判定する問題です。

例えば、nums = [4, 2, 6, 5, 1, 6, 3]、k = 3 の場合、出力は True になります。[6, 3]、[6, 2, 1]、[4, 5] のように分割すると、それぞれの合計がすべて 9 で等しくなるためです。

解決のアプローチ

この問題は、深さ優先探索（DFS）によって各要素をどの部分集合に割り当てるかを全て試すことで解けます。手順は以下の通りです。

• 関数 check() を定義する。配列 v を受け取り、すべての要素が等しいかどうかを確認する。
• i := 1 で初期化し、i < v のサイズである間、i を1ずつ増やしながら以下を繰り返す。
  • v[i] が v[0] と等しくなければ false を返す。
• true を返す。
• 関数 dfs() を定義する。idx、配列 nums、配列 temp を受け取る。
• idx が nums のサイズと同じ場合、check(temp) の結果を返す。
• ret := false とする。
• i := 0 で初期化し、i < temp のサイズである間、i を1ずつ増やしながら以下を繰り返す。
  • temp[i] := temp[i] + nums[idx]
  • ret := dfs(idx + 1, nums, temp)
  • ret が true であれば true を返す。
  • temp[i] := temp[i] − nums[idx]（バックトラックして状態を元に戻す）
• false を返す。
• main 関数では以下を行う。
  • サイズ k の配列 temp を定義する。
  • dfs(0, nums, temp) の結果を返す。

C++での実装例

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

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   bool check(vector<int>& v) {
      for (int i = 1; i < v.size(); i++) {
         if (v[i] != v[0])
            return false;
      }
      return true;
   }
   bool dfs(int idx, vector<int>& nums, vector<int>& temp) {
      if (idx == nums.size()) {
         return check(temp);
      }
      bool ret = false;
      for (int i = 0; i < temp.size(); i++) {
         temp[i] += nums[idx];
         ret = dfs(idx + 1, nums, temp);
         if (ret)
            return true;
         temp[i] -= nums[idx];
      }
      return false;
   }
   bool solve(vector<int>& nums, int k) {
      vector<int> temp(k);
      return dfs(0, nums, temp);
   }
};
bool solve(vector<int>& nums, int k) {
   return (new Solution())->solve(nums, k);
}
int main(){
   vector<int> v = {4, 2, 6, 5, 1, 6, 3};
   int k = 3;
   cout << solve(v, 3);
}

入力

{4, 2, 6, 5, 1, 6, 3}, 3

出力

1

計算量に関する注意点

この手法は各要素を k 個の部分集合のいずれかに割り当てる組み合わせを全て試すため、最悪の場合の時間計算量は O(k^n)（n は要素数）となり、指数オーダーになります。そのため、小規模な入力には有効ですが、大きなデータセットに対しては、事前に総和が k で割り切れるかを確認したり、枝刈り（合計が目標値を超えた部分集合への割り当てを打ち切るなど）を加えることで効率化を図ることが推奨されます。