C++で解く最大平均部分配列II：二分探索による効率的アプローチ

問題概要

n個の整数からなる配列が与えられたとき、長さがk以上の連続する部分配列の中で平均値が最大になるものを見つけ、その最大平均値を出力してください。

例として、入力が [1, 12, -5, -6, 50, 3]、k = 4 の場合を考えてみます。

• 長さ4の場合の最大平均値：12.75（例：[12, -5, -6, 50] の和は51、51 ÷ 4 = 12.75）
• 長さ5の場合の最大平均値：10.8（例：[12, -5, -6, 50, 3] の和は54、54 ÷ 5 = 10.8）
• 長さ6の場合の最大平均値：9.16667（配列全体の和は55、55 ÷ 6 ≒ 9.16667）

これらを比較すると長さ4のときが最も大きいため、答えは 12.75 となります。

解法の考え方：答えに対する二分探索

この問題は、すべての部分配列を列挙すると計算量が膨大になってしまいます。そこで有効なのが「答えそのものを二分探索する」テクニックです。

「平均値が x 以上になるような長さ k 以上の部分配列が存在するか？」を判定する関数 ok(x) を用意すれば、x が大きすぎると false、小さければ true となる単調性が成り立ちます。この性質を利用して x の範囲を絞り込んでいくのがポイントです。

判定関数 ok(x, nums, k) の手順

1. n := nums のサイズとします。
2. サイズ n の配列 arr を定義します。
3. i = 0 から n-1 までの各 i について、arr[i] := nums[i] − x とします。
4. sum := 0、last := 0 で初期化します。
5. i = 0 から k-1 まで sum += arr[i] を実行し、先頭 k 個の和を求めます。
6. sum ≥ 0 であれば true を返します。
7. i = 0、j = k から開始し、j < n の間、i と j を同時に 1 ずつ進めながら以下を繰り返します。
   • last += arr[i]
   • sum += arr[j]
   • last < 0 の場合は sum -= last として、last := 0 にリセットします。
   • sum ≥ 0 であれば true を返します。
8. 条件を満たす区間が一度も現れなければ false を返します。

この判定の鍵は、各要素から x を引いた配列 arr を作る点にあります。「平均が x 以上の部分配列が存在する」ことは「和が 0 以上の長さ k 以上の部分配列が存在する」ことと完全に等価だからです。変数 last は途中までの負の寄与を記録しておき、それが足かせになっている場合は切り捨てることで、長さ k 以上のあらゆる区間を見逃さずチェックできます。

メイン処理：実数上の二分探索

本体側では次のように進めます。

1. ret := 0、low := −∞、high := +∞ で初期化します。
2. high − low > 10⁻⁵ の間、以下を繰り返します。
   • mid := low + (high − low) / 2
   • ok(mid, nums, k) が true なら low := mid とし、ret := mid を更新します。
   • false なら high := mid とします。
3. ループ終了後、ret を返します。

C++ 実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    bool ok(double x, vector <int>& nums, int k){
        int n = nums.size();
        double arr[n];
        for (int i = 0; i < n; i++) {
            arr[i] = nums[i] - x;
        }
        double sum = 0;
        double last = 0;
        for (int i = 0; i < k; i++) {
            sum += arr[i];
        }
        if (sum >= 0)
        return true;
        for (int i = 0, j = k; j < n; i++, j++) {
            last += arr[i];
            sum += arr[j];
            if (last < 0) {
                sum -= last;
                last = 0;
            }
            if (sum >= 0)
            return true;
        }
        return false;
    }
    double findMaxAverage(vector<int>& nums, int k) {
        double ret = 0;
        double low = INT_MIN;
        double high = INT_MAX;
        while (high - low > 1e-5) {
            double mid = low + (high - low) / 2;
            if (ok(mid, nums, k)) {
                low = mid;
                ret = mid;
            } else {
                high = mid;
            }
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,12,-5,-6,50,3};
    cout << (ob.findMaxAverage(v, 4));
}

入力

{1,12,-5,-6,50,3},4

出力

12.75000

計算量の目安

判定関数 ok() は配列を一度走査するだけなので O(n) です。二分探索は精度 10⁻⁵ に収束するまで繰り返されますが、探索範囲が32ビット整数程度であれば反復回数はおよそ60回前後に収まります。したがって全体の計算量は O(n log((high−low)/ε)) となり、すべての候補区間を調べる O(n²) の素朴な解法に比べて大幅に高速です。