【C++】そのサイズのすべてのサブ配列の合計がk以下となる最大サブ配列サイズの求め方

このチュートリアルでは、「そのサイズのすべてのサブ配列(部分配列)の合計がk未満となる、最大のサブ配列サイズ」を求めるプログラムについて解説します。

サイズNの配列と整数kが与えられます。求めるのは、与えられた配列内のその長さのすべてのサブ配列の合計がk以下となるようなサブ配列の長さです。

アルゴリズムの考え方

この問題は「累積和(プレフィックスサム)」と「二分探索」を組み合わせることで、O(N log N)の計算量で効率的に解くことができます。

解法の手順

1. 配列の累積和を事前に計算します。これにより、任意の区間の合計をO(1)で取得できます。
2. サブ配列の長さを候補として二分探索を行います。
3. 各候補の長さmidについて、その長さのすべてのサブ配列の合計がk以下であるかを累積和でチェックします。
4. 条件を満たす場合はより長い長さを、満たさない場合はより短い長さを探索します。

C++での実装例

#include<bits/stdc++.h>
using namespace std;
// 最大長のサブ配列を二分探索で求める
int bsearch(int prefixsum[], int n, int k) {
    int ans = -1;
    // 二分探索の実行
    int left = 1, right = n;
    while (left <= right) {
        int mid = (left + right) / 2;
        int i;
        for (i = mid; i <= n; i++) {
            if (prefixsum[i] - prefixsum[i - mid] > k)
                break;
        }
        if (i == n + 1) {
            left = mid + 1;
            ans = mid;
        }
        else right = mid - 1;
    }
    return ans;
}
// 最大サブ配列サイズを返す
int maxSize(int arr[], int n, int k) {
    int prefixsum[n + 1];
    memset(prefixsum, 0, sizeof(prefixsum));
    for (int i = 0; i < n; i++)
        prefixsum[i + 1] = prefixsum[i] + arr[i];
    return bsearch(prefixsum, n, k);
}
int main() {
    int arr[] = {1, 2, 10, 4};
    int n = sizeof(arr) / sizeof(arr[0]);
    int k = 14;
    cout << maxSize(arr, n, k) << endl;
    return 0;
}

出力

2

コードの解説

bsearch関数:二分探索により、条件を満たす最大のサブ配列長を求めます。長さmidのすべてのサブ配列がk以下であれば答えを更新し、より長い長さを探索します。

maxSize関数:累積和配列prefixsumを構築し、bsearchを呼び出します。prefixsum[i] - prefixsum[i - mid]を計算することで、長さmidのサブ配列の合計を即座に求められます。

実行結果の例:配列 {1, 2, 10, 4}、k = 14 の場合を考えます。長さ2のサブ配列 {1, 2}、{2, 10}、{10, 4} の合計はそれぞれ3、12、14であり、すべてk以下です。一方、長さ3のサブ配列 {2, 10, 4} の合計は16となりkを超えるため、答えは2になります。