Count of Range Sum

Count of Range Sum: Given an integer array nums and two integers lower and upper, 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 inclusive, where i <= j.

Example 1:

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

Example 2:

Input: nums = [0], lower = 0, upper = 0
Output: 1

Constraints:

1 <= nums.length <= 10^5
-2^31 <= nums[i] <= 2^31 - 1
-10^5 <= lower <= upper <= 10^5
The answer is guaranteed to fit in a 32-bit integer.

Try this Problem on your own or check similar problems:

Solution:

Java
JavaScript
Python
C++

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
      long[] sums = new long[nums.length + 1];
      long prefixSum = 0;
      for(int i = 1; i <= nums.length; ++i){
          prefixSum += nums[i - 1];
          sums[i] = prefixSum;
      }
      return helper(sums, lower, upper, 0, sums.length);
  }

  private int helper(long[] sums, int lower, int upper, int start, int end) {
      if (end - start <= 1) return 0;

      int mid = start + (end - start) / 2;
      int countLeft =  helper(sums, lower, upper, start, mid);
      int countRight = helper(sums, lower, upper, mid, end);
      int totalCount = countLeft + countRight;

      int j = mid, k = mid, l = mid;
      long[] temp = new long[end - start];
      int tempWriter = 0;
      for (int i = start; i < mid; ++i) {
          while (j < end && sums[j] - sums[i] <= upper) j++;
          while (k < end && sums[k] - sums[i] < lower) k++;
          while (l < end && sums[l] < sums[i]) temp[tempWriter++] = sums[l++];
          temp[tempWriter++] = sums[i];
          totalCount += j - k;
      }
      System.arraycopy(temp, 0, sums, start, l - start);
      return totalCount;
  }
}

/**
 * @param {number[]} nums
 * @param {number} lower
 * @param {number} upper
 * @return {number}
 */
var countRangeSum = function (nums, lower, upper) {
  const n = nums.length;
  const sums = new Array(n + 1).fill(0);
  let prefixSum = 0;
  for (let i = 1; i <= n; ++i) {
    prefixSum += nums[i - 1];
    sums[i] = prefixSum;
  }
  return helper(sums, lower, upper, 0, sums.length);
};

function helper(sums, lower, upper, start, end) {
  if (end - start <= 1) return 0;

  const mid = start + Math.floor((end - start) / 2);
  let countLeft = helper(sums, lower, upper, start, mid);
  let countRight = helper(sums, lower, upper, mid, end);
  let totalCount = countLeft + countRight;

  let j = mid;
  let k = mid;
  let l = mid;
  const temp = new Array(end - start);
  let tempWriter = 0;

  for (let i = start; i < mid; ++i) {
    while (j < end && sums[j] - sums[i] <= upper) j++;
    while (k < end && sums[k] - sums[i] < lower) k++;
    while (l < end && sums[l] < sums[i]) temp[tempWriter++] = sums[l++];
    temp[tempWriter++] = sums[i];
    totalCount += j - k;
  }

  sums.splice(start, l - start, ...temp.slice(0, l - start));
  return totalCount;
}

class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        n = len(nums)
        sums = [0] * (n + 1)
        prefixSum = 0
        for i in range(1, n + 1):
            prefixSum += nums[i - 1]
            sums[i] = prefixSum
        return self.helper(sums, lower, upper, 0, len(sums))

    def helper(self, sums: List[int], lower: int, upper: int, start: int, end: int) -> int:
        if end - start <= 1:
            return 0

        mid = start + (end - start) // 2
        countLeft = self.helper(sums, lower, upper, start, mid)
        countRight = self.helper(sums, lower, upper, mid, end)
        totalCount = countLeft + countRight

        j, k, l = mid, mid, mid
        temp = [0] * (end - start)
        tempWriter = 0
        for i in range(start, mid):
            while j < end and sums[j] - sums[i] <= upper:
                j += 1
            while k < end and sums[k] - sums[i] < lower:
                k += 1
            while l < end and sums[l] < sums[i]:
                temp[tempWriter] = sums[l]
                tempWriter += 1
                l += 1
            temp[tempWriter] = sums[i]
            tempWriter += 1
            totalCount += j - k

        sums[start:start + (l - start)] = temp[:l - start]
        return totalCount

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size();
        vector<long long> sums(n + 1);
        long long prefixSum = 0;
        for (int i = 1; i <= n; ++i) {
            prefixSum += nums[i - 1];
            sums[i] = prefixSum;
        }
        return helper(sums, lower, upper, 0, sums.size());
    }

private:
    int helper(vector<long long>& sums, int lower, int upper, int start, int end) {
        if (end - start <= 1)
            return 0;

        int mid = start + (end - start) / 2;
        int countLeft = helper(sums, lower, upper, start, mid);
        int countRight = helper(sums, lower, upper, mid, end);
        int totalCount = countLeft + countRight;

        int j = mid, k = mid, l = mid;
        vector<long long> temp(end - start);
        int tempWriter = 0;
        for (int i = start; i < mid; ++i) {
            while (j < end && sums[j] - sums[i] <= upper)
                j++;
            while (k < end && sums[k] - sums[i] < lower)
                k++;
            while (l < end && sums[l] < sums[i])
                temp[tempWriter++] = sums[l++];
            temp[tempWriter++] = sums[i];
            totalCount += j - k;
        }
        copy(temp.begin(), temp.begin() + (l - start), sums.begin() + start);
        return totalCount;
    }
};

Time/Space Complexity:

Time Complexity: O(nlogn)
Space Complexity: O(n)

Explanation:

It's a more advanced two pointers approach where we do a merge sort, but during the merge phase we do additional work like counting. Let's break down the implementation:

First since we're interested in subarray (range) sums, it's natural to utilize the prefix sum like we did in Range Query Sum Immutable. We will be doing the merge sort and counting on the prefixsum array, not the original input.
We set up regular merge sort, in which we half the array we're sorting, do a merge sort on both halves and merge them together using something like Merge Two Sorted Lists (same approach can be used on arrays). In the merge phase we keep the temp where we will store the sorted order of two halves, and after we're done with counting and merging we copy it to the prefix array using System.arraycopy(temp, 0, sums, start, l - start); (copy from temp source from index 0, to destination sums from index start total of l-start elements).
While we're merging the two halves we also count the number of sums in lower..upper range using the prefixSum[a]-prefixSum[b] approach to get a sum between indices a..b. Since the second half is also sorted, we use the two pointers approach and move the pointer j until we have invalid sum (larger than upper), this will just offset the total count by one (instead of doing j-k+1), and we move the pointer k until we reach the first valid sum larger or equal to lower.
totalCount is sum of counts from left and right side and cross sums between left and right side.

The time complexity is that of merge sort O(nlogn) while we also have linear space complexity because of the sums prefix sum array and the temp for merging purposes. Why don't we care about the relative order of sums, and can utilize order (sorted array) to implement the two pointers approach? Remember that we're operating on prefix sum array sums, not the original input array and also we're looking only for number of prefix sum pairs whose difference is in the desired boundary (lower, upper).

Example 1:​

Example 2:​

Constraints:​

Solution:​

Time/Space Complexity:​