すべての部分配列の合計がk未満となる最大の部分配列サイズを求めるC++プログラム

この記事では、n個の正の整数からなる配列 arr[] と整数 k が与えられたときに、「そのサイズのすべての部分配列の合計が k 未満となる最大の部分配列サイズ」を求めるC++プログラムを紹介します。

問題の概要

配列の要素から構成できるあるサイズの部分配列をすべて考えたとき、どの部分配列を取っても要素の合計が k 以下になるような、最大のサイズを見つけるのが目標です。

具体例で理解しよう

入力

arr[n] = {4, 1, 3, 2}, k = 9

出力

3

説明

サイズ3の部分配列とその合計は次のようになります。

{4, 1, 3} = 8
{1, 3, 2} = 6

サイズ3のすべての部分配列の合計が k(=9)以下であるため、答えは 3 となります。

解法1:累積和(プレフィックスサム)と二分探索

まず基本となる解法として、累積和(prefix sum)を利用する方法があります。累積和とは、配列の先頭から各インデックスまでの要素の合計を記録した配列のことです。

この累積和に対して二分探索を行い、合計が k を超えない最大の長さを求めます。あるサイズで合計が k を超えてしまうなら、それより大きいサイズでは必ず条件を満たさないという性質があるため、二分探索によって効率よく境界を見つけられます。

実装例

#include<iostream>
using namespace std;

int calcSubArraySize(int arr[], int n, int k){
    int prefixSum[n + 1];
    prefixSum[0] = 0;
    for (int i = 0; i < n; i++)
        prefixSum[i + 1] = prefixSum[i] + arr[i];

    // 条件を満たすサイズを二分探索で求める
    int maxLen = -1;
    int start = 1, end = n;
    while (start <= end){
        int mid = (start + end) / 2;
        int i;
        for (i = mid; i <= n; i++){
            if (prefixSum[i] - prefixSum[i - mid] > k)
                break;
        }
        if (i == n + 1){
            start = mid + 1;
            maxLen = mid;
        }
        else
            end = mid - 1;
    }
    return maxLen;
}

int main(){
    int arr[] = {4, 1, 2, 3};
    int n = sizeof(arr) / sizeof(arr[0]);
    int k = 9;
    cout<<"The maximum subarray size, such that all subarrays of that size have sum less than k is "<<calcSubArraySize(arr, n, k);
    return 0;
}

この方法は十分に効率的ですが、さらに優れたアプローチが存在します。

解法2:スライディングウィンドウ法

より洗練された方法として、スライディングウィンドウ(sliding window)法を使って部分配列の合計を管理する手法があります。

まず配列全体をウィンドウとして合計を計算し、合計が k を超えたら左端から要素を外していき、再び k 以下になるまでウィンドウを縮めます。この操作で得られる最小のウィンドウ長から 1 を引いた値が、「すべての部分配列の合計が k 以下となる最大サイズ」になります。

実装例

#include <iostream>
using namespace std;

int calcSubArraySizeSW(int arr[], int n, int k){
    int maxLen = n;
    int subArraySum = 0;
    int start = 0;
    for (int end = 0; end < n; end++){
        subArraySum += arr[end];
        while (subArraySum > k) {
            subArraySum -= arr[start];
            start++;
            maxLen = min(maxLen, end - start + 1);
            if (subArraySum == 0)
                break;
        }
        if (subArraySum == 0) {
            maxLen = -1;
            break;
        }
    }
    return maxLen;
}

int main(){
    int arr[] = { 4, 1, 3, 2, 6 };
    int k = 12;
    int n = sizeof(arr)/ sizeof(arr[0]);
    cout<<"The maximum subarray size, such that all subarrays of that size have sum less than k is "<<calcSubArraySizeSW(arr, n, k);
    return 0;
}

出力

The maximum subarray size, such that all subarrays of that size have sum less than k is 4

(訳:すべての部分配列の合計が k 未満となる最大の部分配列サイズは 4 です)

まとめ

累積和と二分探索を組み合わせた解法でも O(n log n) で求解できますが、スライディングウィンドウ法を使えば O(n) の計算量でより効率的に答えを得られます。なお、どちらの解法も配列が正の整数のみで構成されていることを前提としている点に注意してください。