327. Count of Range Sum

[huimeich]

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: Given nums = [-2, 5, -1], lower = -2, upper = 2, Return 3. The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

[huimeich]

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        if (nums.empty()) return 0;
        vector<long> sum;    sum.push_back(0);
        for (int i=0; i<nums.size(); i++) sum.push_back(nums[i]+sum[i]);
        return countR(sum, 0, nums.size()+1, lower, upper);
    }
    int countR(vector<long>& sum, int left, int right, int lower, int upper) {
        if (lower>upper || right-left<=1) return 0;
        int count;
        int mid=left+(right-left)/2; int i, k=mid, l=mid, r=0, j=mid;
        count=countR(sum, left, mid, lower, upper)+countR(sum, mid, right, lower, upper);
        vector<long> cache(right-left,0); 
        for (i=left; i<mid; i++) {
            while (l<right && sum[l]-sum[i]<lower) l++;
            while (j<right && sum[j]-sum[i]<=upper) j++;
            count+=j-l;
            while (k<right && sum[k]<sum[i]) {cache[r++]=(sum[k++]);}
            cache[r++]=(sum[i]);
        }
        for (i=0; i<k-left; i++) sum[i+left]=cache[i];
        return count;
    }
};