題目網址：https://leetcode.cn/problems/count-of-range-sum/

題意：給一整數 array nums 以及兩個整數 lower 和 upper。求 nums 中, 值位於 [lower, upper] 之區間和個數。

區間和 S(i, j) 表示在 nums 中, idx 從 i 到 j 中的區間之和, 包含 i 和 j（i ≤ j）。

Solution：

想法：利用 Prefix Sum + Binary Search + Merge Sort, 類似 303. Range Sum Query - Immutable。presum[0] 必須為 0, 若我們想求一區間和 S(i, j) 可看做是 prefix[j + 1] - prefix[i]

我們先考慮以下的問題:

給兩個升序的 array n1 和 n2, 找出所有 (i, j) 滿足 : n2[j] - n2[i] ∈ [lower,upper]

  • 我們在 n2 中使用兩個 ptr left 和 right, 初始皆為 n2 的起始位置。
  • 考慮 n1 的第一個元素
    • 將 left 不斷右移, 直到 n2[left] ≥ n1[0] + lower 為止
      ➔ 此時, left 以及其右邊元素皆 ≥ n1[0] + lower
    • 將 right 不斷右移, 直到 n2[right] > n1[0] + upper 為止
      ➔ 此時, right 的左邊元素皆 ≤ n1[0] + upper
    • 得到區間 [left, right) 中所有的 j 皆滿足 : n2[j] - n1[0] ∈ [lower,upper]
  • 接著考慮 n1 的第二個元素, 依此類推, 遍歷完 n1 後可得到所有符合條件的 (i, j) 之數量

但 presum 不保證是升序的, 因為 nums[i] 有可能為負數, 因此採用 Merge Sort 來解決此問題

  • 首先, 我們將 presum 的元素都拆成一個一組, 然後開始兩兩合併
  • 合併的過程中, 會先遍歷左區間所有的idx i, 然後到右區間去找符合條件的 idx j 之個數
    找到後, 合併兩區間, 並由小到大排序, 將問題轉換成上面討論的問題

e.g. nums = [-2, 5, -1] ➔ presum = [0, -2, 3, 2]

  • [0]、[-2] 合併時, 左區間元素 = 0, 右區間 x = 1, y = 2 ➔ 得到一組解
  • [3]、[2] 合併時, 左區間元素 = 3, 右區間 x = 3, y = 4 ➔ 得到一組解
  • [-2, 0]、[2, 3] 合併時
    • 左區間第一個元素 = -2, 右區間 x = 2, y = 2 ➔ 沒有滿足條件之 j
    • 左區間第一個元素 = 0, 右區間 x = 2, y = 3 ➔ 得到一組解

class Solution {
private:
    int res = 0;

    void mergeSort(vector<long>& presum, int left, int right, int lower, int upper){
        if (left >= right) {
            return;
        }

        // 將 [left, right] 拆成左、右兩個區間, 分別是 [left, mid] 和 [mid + 1, right]
        const int mid = left + (right - left) / 2;
        mergeSort(presum, left, mid, lower, upper);
        mergeSort(presum, mid + 1, right, lower, upper);

        int x = mid + 1, y = mid + 1;

        // 遍歷左區間的元素, 並從右區間中找出滿足:
        //    prefix[j] - prefix[i] in [lower, upper] 的個數, 其中 x ≤ j < y
        // 代表 prefix[i] 可找到 (y-x) 個 j 滿足題目條件
        for (int i = left; i <= mid; ++i) {
            // 從右區間中找到第一個 >= lower 的 idx
            while (x <= right && presum[x] - presum[i] < lower) {
                ++x;
            }

            // 從右區間中找到第一個 > upper 的 idx, 則 y - 1 為 <= upper 之 idx
            while (y <= right && presum[y] - presum[i] <= upper) {
                ++y;
            }

            // 區間 [x, y) 的個數, 代表滿足區間和 [lower, upper] 之個數
            res += (y - x);
        }

        // 對左、右區間進行 merge 並 sorting
        merge(presum, left, mid, right);
    }

    // 對左、右區間進行 merge 並 sorting, 只需 O(n) time
    void merge(vector<long>& presum, int& left, int& mid, int& right){
        int p = 0;
        int p1 = left, p2 = mid + 1;
        vector<long> sorted(right - left + 1);

        // 跳出 loop 代表其中一個區間已做完
        while (p1 <= mid && p2 <= right) {
            if (presum[p1] < presum[p2]) {
                sorted[p++] = presum[p1++];
            } else {
                sorted[p++] = presum[p2++];
            }
        }

        // 左區間已做完, 將右區間剩下的元素加入到 sorted 中
        if (p1 > mid) {
            for (int i = p2; i <= right; ++i) {
                sorted[p++] = presum[i];
            }
        } else { // 右區間已做完, 將左區間剩下的元素加入到 sorted 中
            for (int i = p1; i <= mid; ++i) {
                sorted[p++] = presum[i];
            }
        }

        // 用 sorted 覆蓋掉原本 presum 的 [left, right) 區間中未排序的值
        for (int i = 0; i < sorted.size(); ++i) {
            presum[left + i] = sorted[i];
        }
    }

public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        const int n = nums.size();
        vector<long> presum(n+1, 0);

        for (int i = 0; i < n; ++i) {
            presum[i+1] = presum[i] + nums[i];
        }

        mergeSort(presum, 0, n, lower, upper);

        return res;
    }
};
  • time：$O(n \cdot log(n))$ ➔ $T(n) = 2 \cdot T(\dfrac{n}{2}) + O(n)$ 利用 Master theorem 計算得到
    • $2 \cdot T(\dfrac{n}{2})$：遞迴左、右區間
    • $O(n)$：merge 左、右區間並做 sorting
  • space：$O(n)$ ➔ $O(log(n)) + O(n)$
    • $O(log(n))$：遞迴深度
    • $O(n)$：每層遞迴產生的 sorted 長度不超過 n