【C++】範囲 [lower, upper] に含まれる区間和の個数を効率的に数える方法

問題概要

整数配列 nums が与えられたとき、範囲 [lower, upper]（両端を含む）に収まる区間和の個数を求めます。ここで区間和 S(i, j) とは、i ≤ j を満たすインデックス i から j までの nums の要素の総和として定義されます。

例えば、入力が [-3, 6, -1]、lower = -2、upper = 2 の場合、答えは 2 になります。条件を満たすのは [0, 2]（合計 2）と [2, 2]（合計 -2）の2つだけだからです。

解法のアプローチ：累積和とマージソートの組み合わせ

この問題は素朴に全ペアを調べると O(n²) かかりますが、累積和（プレフィックスサム）とマージソートを組み合わせることで O(n log n) まで高速化できます。

ポイントは、区間和 S(i, j) = prefix[j] − prefix[i] と書き直せることです。「lower ≤ prefix[j] − prefix[i] ≤ upper」となるペア (i, j) の個数を数えればよく、ソート済みの配列上では two-pointer テクニックで各 i に対する該当数を一度に求められます。

アルゴリズムの手順

• 関数 mergeIt() を定義します。引数は配列 prefix、start、mid、end、lower、upper です。
• i := start、j := mid + 1
• temp := end − start + 1
• low := mid + 1、high := mid + 1
• k := 0
• サイズ temp の作業用配列 arr を用意します。
• i ≤ mid の間、以下を繰り返します。
  • low ≤ end かつ prefix[low] − prefix[i] < lower の間、low を1ずつ増やします。
  • high ≤ end かつ prefix[high] − prefix[i] ≤ upper の間、high を1ずつ増やします。
  • j ≤ end かつ prefix[j] < prefix[i] の間、arr[k] := prefix[j] として j と k を1ずつ増やします。
  • arr[k] := prefix[i] として i と k を1ずつ増やします。
  • count := count + high − low（この i を始点とする有効な区間の数を加算）
• j ≤ end の間、残りの要素を arr にコピーします。
• 最後に arr の内容を prefix[start] 以降へ書き戻してソート済み状態を維持します。

再帰的なマージ処理

• 関数 merge(prefix[], start, end, lower, upper) を定義します。
• start ≥ end なら何もせず return します。
• mid := start + (end − start) / 2
• merge(prefix, start, mid, lower, upper) を呼び出します。
• merge(prefix, mid + 1, end, lower, upper) を呼び出します。
• mergeIt(prefix, start, mid, end, lower, upper) を呼び出して統合しながらカウントします。

メイン処理

• n := nums のサイズ、count := 0 と初期化します。
• サイズ n+1 の累積和配列 prefix を用意し、prefix[0] := 0 とします。
• i = 1 から n まで、prefix[i] := prefix[i−1] + nums[i−1] を計算します。
• merge(prefix, 0, n, lower, upper) を呼び出し、最後に count を返します。

C++ 実装例

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
public:
    int count = 0;
    void mergeIt(lli prefix[], lli start ,lli mid, lli end, lli lower, lli upper){
        lli i = start, j = mid + 1;
        lli temp = end - start + 1;
        lli low = mid + 1, high = mid + 1;
        lli k = 0;
        lli arr[temp];
        while(i <= mid){
            while(low <= end && prefix[low] - prefix[i] < lower) low++;
            while(high <= end && prefix[high] - prefix[i] <= upper) high++;
            while(j<= end && prefix[j] < prefix[i]){
                arr[k] = prefix[j];
                j++;
                k++;
            }
            arr[k] = prefix[i];
            i++;
            k++;
            count += high - low;
        }
        while(j <= end){
            arr[k] = prefix[j];
            k++;
            j++;
        }
        for(i = 0; i < temp; i++){
            prefix[start] = arr[i];
            start++;
        }
    }
    void merge(lli prefix[], lli start, lli end, lli lower, lli upper){
        if(start >= end)return;
        lli mid = start + (end - start) / 2;
        merge(prefix, start, mid, lower, upper);
        merge(prefix, mid + 1, end, lower, upper);
        mergeIt(prefix, start, mid, end, lower, upper);
    }
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        lli n = nums.size();
        count = 0;
        lli prefix[n + 1];
        prefix[0] = 0;
        for(lli i = 1; i <= n; i++){
            prefix[i] = prefix[i - 1] + nums[i - 1];
        }
        merge(prefix, 0, n, lower, upper);
        return count;
    }
};
main(){
    Solution ob;
    vector<int> v = {-3,6,-1};
    cout << (ob.countRangeSum(v, -2, 2));
}

入力

{-3,6,-1}
-2
2

出力

2

まとめ

累積和によって「区間和の個数」問題を「差が範囲に収まるペアの個数」問題に変換できるのが本手法の核心です。マージソートの統合フェーズで two-pointer を使うことで、各要素ごとの探索が償却的に高速になり、全体の計算量は O(n log n) に抑えられます。負の値を含む配列でも正しく動作する点も、このアプローチの大きな利点です。