327. Count of Range Sum

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive. Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.

Example:

Input: nums = [-2,5,-1], lower = -2, upper = 2,
Output: 3 
Explanation: The three ranges are : [0,0], [2,2], [0,2] and their respective sums are: -2, -1, 2.

Solution

Approach 1: Merge sort.

Code (Python)

Approach 1:

Code (C++)

Approach 1:

class Solution {
private:
    int count;
    void doMergeSort(vector<long>& sums, int head, int tail, int lower, int upper) {
        if (head >= tail)
            return;
        int mid = head + (tail - head) / 2;
        doMergeSort(sums, head, mid, lower, upper);
        doMergeSort(sums, mid+1, tail, lower, upper);
        vector<int> tmp;
        int i1 = head;
        int i2 = mid + 1;
        int iLower2 = mid + 1;
        int iUpper2 = mid + 1;
        while (i1 <= mid || i2 <= tail)
        {
            if (i2 > tail || i1 <= mid && sums[i1] <= sums[i2])
            {
                while (iLower2 <= tail && sums[iLower2] - sums[i1] < lower)
                    iLower2++;
                while (iUpper2 <= tail && sums[iUpper2] - sums[i1] <= upper)
                    iUpper2++;
                count += iUpper2 - iLower2;
                tmp.push_back(sums[i1]);
                i1++;
            }
            else
            {
                tmp.push_back(sums[i2]);
                i2++;
            }
        }
        for (int i = head; i <= tail; ++i)
        {
            sums[i] = tmp[i-head];
        }
    }
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size();
        if (n == 0)
            return 0;
        count = 0;
        vector<long> sums;
        sums.push_back(0);
        for (int i = 0; i < n; ++i)
            sums.push_back(sums[i] + nums[i]);
        doMergeSort(sums, 0, sums.size()-1, lower, upper);
        return count;
    }
};