C++で同じ平均値を持つ2つのグループに配列を分割できるか判定する方法

配列 A が与えられたとします。この配列のすべての要素を、リスト B またはリスト C のいずれかに振り分ける必要があります(初期状態では B と C はともに空です)。振り分け後、B と C がどちらも空でないという条件のもとで、両者の平均値が等しくなるような振り分けが存在するかどうかを判定します。

たとえば、入力が [1,2,3,4,5,6,7,8,9,10] の場合、答えは true になります。

解法の考え方

まず、この問題の背後にある数学的な性質を整理しましょう。配列全体の要素数を n、合計を S とします。あるグループが k 個の要素を持ち、その合計が s であるとき、両グループの平均が等しくなる条件は次のように表せます。

s / k = (S − s) / (n − k)

これを変形すると s × n = S × k となります。つまり、「k 個の要素を選んだとき、その合計が S × k / n に一致するような組み合わせが存在するか」を調べればよいことになります。

アルゴリズムの手順

• n を配列 A のサイズ、total を 0 として初期化します。
• i を 0 から n−1 まで動かしながら total に各要素を加算し、配列の総和を求めます。
• isPossible を false、m を n / 2 に設定します。
• i を 1 から m まで動かし、(total * i) % n == 0 となる i が存在すれば isPossible を true にします。これは、選ぶ要素数 k に対して目標となる合計 S × k / n が整数になり得るかどうかの事前チェックです。
• isPossible が false のままの場合、条件を満たす分割は不可能なので false を返します。
• サイズ (total + 1) × ((n / 2) + 1) の2次元DPテーブル dp を定義します。dp[j][l] は「l 個の要素を選んで合計 j を作れるかどうか」を表します。
• dp[0][0] = true と初期化します。
• 各要素 x について、j を total から x まで逆順に、l を 1 から n/2 まで動かしながら dp[j][l] = dp[j][l] || dp[j - x][l - 1] によりテーブルを更新します。
• 最後に i を 1 から n/2 まで動かし、(total * i) % n == 0 かつ dp[total * i / n][i] が true であれば true を返します。
• 該当するものが見つからなければ false を返します。

以下の実装例を見ると、より理解が深まるでしょう。

実装例(C++)

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   bool splitArraySameAverage(vector<int>& A) {
      int n = A.size();
      int total = 0 ;
      for(int i = 0; i < n; i++) total += A[i];
      bool isPossible = false;
      int m = n / 2;
      for (int i = 1; i <= m && !isPossible; ++i)
      if (total*i%n == 0) isPossible = true;
      if (!isPossible) return false;
      vector < vector <bool> > dp(total + 1, vector <bool>((n / 2) + 1));
      dp[0][0] = true;
      for(int i = 0; i < n; i++){
         int x = A[i];
         for(int j = total; j >= x; j--){
            for(int l = 1; l <= (n / 2); l++){
               dp[j][l] = dp[j][l] || dp[j - x][l - 1];
            }
         }
      }
      for(int i = 1 ; i <= (n / 2); i++){
         if((total * i) % n == 0 && dp[total * i / n][i]) return true;
      }
      return false;
   }
};
main(){
   Solution ob;
   vector<int> v = {1,2,3,4,5,6,7,8,9,10};
   cout << (ob.splitArraySameAverage(v));
}

入力

{1,2,3,4,5,6,7,8,9,10}

出力

1