C++で配列を最大K個に分割して平均の合計を最大化する方法

問題概要

数値の配列 A が与えられます。この配列を最大 K 個の隣接する(空でない)グループに分割し、スコアを「各グループの平均値の合計」と定義します。このとき、達成できる最大スコアを求めるのが本問題です。

入力例

入力配列が {9, 2, 5, 3, 10} の場合、たとえば次のように分割できます。

{9} {2, 5, 3} {10}

このときの平均の合計は次のとおりです。

9 + (2 + 5 + 3) / 3 + 10 = 22.33

アルゴリズム(メモ化再帰)

この問題は、メモ化(記憶化)再帰を使うことで効率よく解くことができます。

• memo[i][k]:A[i]〜A[n-1] を最大 K 個の部分に分割したときの最高スコアを表します。
• 最初のグループとして A[i]〜A[j-1] を切り出し、残りの A[j]〜A[n-1] をさらに分割することを考えます。このときの候補スコアは average(i, j) + score(j, k-1) です。ここで average(i, j) = (A[i] + A[i+1] + … + A[j-1]) / (j - i) です。
• 考えられるすべての分割位置 j を試し、その中で最も高いスコアを採用します。
• 一般形の再帰式は次のように表せます。
  memo[n][k] = max(memo[n][k], score(i, arr, k-1) + average(i, j))

C++での実装例

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

double memo[MAX][MAX];

double score(int n, vector<int>& arr, int k) {
    if (memo[n][k] > 0) {
        return memo[n][k];
    }
    double sum = 0;
    for (int i = n - 1; i > 0; i--) {
        sum += arr[i];
        memo[n][k] = max(memo[n][k], score(i, arr, k - 1) + sum / (n - i));
    }
    return memo[n][k];
}

double getLargestSum(vector<int>& arr, int K) {
    int n = arr.size();
    double sum = 0;
    memset(memo, 0.0, sizeof(memo));
    for (int i = 0; i < n; i++) {
        sum += arr[i];
        memo[i + 1][1] = sum / (i + 1);
    }
    return score(n, arr, K);
}

int main() {
    vector<int> arr = {9, 2, 5, 3, 10};
    int K = 3;
    cout << "Largest sum = " << getLargestSum(arr, K) << endl;
    return 0;
}

出力

上記のプログラムをコンパイルして実行すると、次の出力が得られます。

Largest sum = 22.3333

計算量

状態の総数は n × K、各状態で分割位置を最大 n 通り試すため、時間計算量は O(n² × K)、空間計算量は O(n × K) となります。また、大きな値ほど単独のグループに分離したほうが平均の合計が大きくなりやすい、という性質を意識すると、この問題の構造がより理解しやすくなります。