Đếm tổng đoạn trong khoảng [lower, upper] trong C++

Cho mảng nums, đếm số cặp (i, j) (i ≤ j) sao cho tổng đoạn S(i, j) nằm trong [lower, upper].

Ví dụ

nums = [-3, 6, -1], lower = -2, upper = 2

Các đoạn hợp lệ: [0,2] (tổng 2), [2,2] (tổng -2) → Kết quả: 2.

Thuật toán: Merge Sort + Prefix Sum

Biến đổi: S(i, j) = prefix[j+1] - prefix[i]. Cần đếm cặp (i, j) với và lower ≤ prefix[j] - prefix[i] ≤ upper.

Sử dụng Merge Sort biến đổi: khi merge 2 nửa đã sắp xếp, dùng 2 con trỏ low, high để đếm số j thỏa mãn cho mỗi i.

Chi tiết mergeIt

• Mảng prefix đã được chia 2 nửa [start..mid] và [mid+1..end] đã sắp xếp.
• Với mỗi i trong nửa trái:
  • Tăng low cho đến khi prefix[low] - prefix[i] ≥ lower.
  • Tăng high cho đến khi prefix[high] - prefix[i] > upper.
  • Số j hợp lệ = high - low.
• Merge 2 nửa về mảng tạm, copy ngược lại.

Cài đặt 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++];
            }
            arr[k++] = prefix[i++];
            count += high - low;
        }
        while (j <= end) arr[k++] = prefix[j++];
        for (i = 0; i < temp; i++) prefix[start++] = arr[i];
    }

    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);
    return 0;
}

Độ phức tạp

• Thời gian: O(n log n) — Merge Sort.
• Không gian: O(n) — mảng prefix và mảng tạm.

Input / Output

Input:  {-3,6,-1} -2 2
Output: 2