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. Build an array of cumulative sum, and convert the problem to find pairs of differences within items of cumulative sum. Recursively sort the items while returning counts.

Code (Python)

Approach 1:

    def countRangeSum(self, nums: list[int], lower: int, upper: int) -> int:
        # O(NlogN) idea: mergesort
        # build array of cumulative sum, cum_sum can be sorted since we're only interested in the differences of their items
        # we want subarrays to be sorted so that it's O(N) to find number of diffs between items in left and item in right
        # return counts in subarrays while mergesorting
        if not nums:
            return 0
        if len(nums) == 1:
            return 1 if lower <= nums[0] <= upper else 0
        cum_sum = [0] # len(cum_sum) == n + 1 to easily handle edge case
        for num in nums:
            cum_sum.append(num + cum_sum[-1])
        # includes 0 but not includes len(cum_sum)
        result = self._sort_and_count(cum_sum, 0, len(cum_sum), lower, upper)
        return result
    
    def _sort_and_count(self, cum_sum, lo, hi, lower, upper):
        mid = (lo + hi) // 2
        if mid == lo: # only 1 item, meaning already sorted
            return 0
        
        # left side does not include mid, right side does
        count = self._sort_and_count(cum_sum, lo, mid, lower, upper) + self._sort_and_count(cum_sum, mid, hi, lower, upper)
        i = mid
        j = mid
        for left in cum_sum[lo:mid]: # mid not included here
            # note that i and j keep their position (not reset to mid each time)
            # meaning it's still a linear operation even with 2 loops
            while i < hi and cum_sum[i] - left < lower:
                i += 1
            while j < hi and cum_sum[j] - left <= upper:
                j += 1
            count += (j - i)
        # assuming sorted() takes linear time to merge to sorted arrays
        cum_sum[lo:hi] = sorted(cum_sum[lo:hi])
        return count

    # O(N^2) (Time Limit Exceeded): build array of cumulative sum, do pairwise difference
    def countRangeSumSlow(self, nums: list[int], lower: int, upper: int) -> int:
        if not nums:
            return 0
        if len(nums) == 1:
            return 1 if lower <= nums[0] <= upper else 0
        cum_sum = [nums[0]]
        for i in range(1, len(nums)):
            cum_sum.append(nums[i] + cum_sum[-1])
        totals = 0
        for i in range(len(nums)):
            for j in range(i, len(nums)):
                if i == 0:
                    before = 0
                else:
                    before = cum_sum[i-1]
                if lower <= cum_sum[j] - before <= upper:
                    totals += 1
        return totals

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;
    }
};