นับผลรวมช่วงใน C++ (Count of Range Sum)

ภาพรวมปัญหา

ให้อาเรย์ nums หาจำนวนช่วง (i, j) ที่ i ≤ j และผลรวม nums[i..j] อยู่ในช่วง [lower, upper]

ตัวอย่าง

nums = [-3,6,-1], lower = -2, upper = 2 → Output: 2
ช่วง: [0,2] sum=2, [2,2] sum=-2

วิธีแก้ปัญหา: Merge Sort + Prefix Sum

ใช้ prefix sum: prefix[i] = sum(nums[0..i-1])
ผลรวมช่วง [i,j] = prefix[j+1] - prefix[i]
ต้องนับคู่ (i,j) ที่ lower ≤ prefix[j+1] - prefix[i] ≤ upper

ทำได้ด้วย Modified Merge Sort: ขณะ merge นับจำนวนคู่ที่ตรงเงื่อนไขได้ใน O(n log n)

ขั้นตอนหลัก

1. สร้าง prefix sum array ขนาด n+1
2. ฟังก์ชัน mergeIt(prefix, start, mid, end, lower, upper):
   • นับคู่ (i,j) ที่ i ∈ [start,mid], j ∈ [mid+1,end] และ lower ≤ prefix[j] - prefix[i] ≤ upper
   • ใช้ two pointers low และ high บนครึ่งขวา (ซึ่งเรียงแล้ว)
   • merge ทั้งสองครึ่งแบบปกติ
3. เรียก merge แบบ divide-and-conquer

โค้ด 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;
    }
};

int main() {
    Solution ob;
    vector<int> v = {-3,6,-1};
    cout << (ob.countRangeSum(v, -2, 2));
}

การวิเคราะห์ความซับซ้อน

• เวลา: O(n log n)
• พื้นที่: O(n)