Skip to main content

Count of Range Sum

Problem

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 <= 105
  • -231 <= nums[i] <= 231 - 1
  • -105 <= lower <= upper <= 105
  • The answer is guaranteed to fit in a 32-bit integer.

Solution

/**
* @param {number[]} nums
* @param {number} lower
* @param {number} upper
* @return {number}
*/
var countRangeSum = function(nums, lower, upper) {
const prefixSum = nums.reduce((acc, cur) => {
acc.push(acc[acc.length - 1] + cur);
return acc;
}, [0]);

const merge = (lo, mid, hi, arr) => {
const lst = [];
let i1 = lo; // index of left half
let i2 = mid + 1; // index of right half

// merge arr[lo:mid + 1] and arr[mid + 1:hi + 1] into lst
while (i1 <= mid || i2 <= hi) {
if (i1 > mid) {
lst.push(arr[i2]);
i2++;
} else if (i2 > hi) {
lst.push(arr[i1]);
i1++;
} else if (arr[i1] < arr[i2]) {
lst.push(arr[i1]);
i1++;
} else {
lst.push(arr[i2]);
i2++;
}
}

// copy lst over into arr
for (let i = lo; i <= hi; i++) {
arr[i] = lst[i - lo];
}
}

const mergeSort = (lo, hi, arr) => {
if (lo >= hi) { // base case
return 0;
}
const mid = Math.floor((lo + hi) / 2); // mid index
let count = mergeSort(lo, mid, arr) + mergeSort(mid + 1, hi, arr); // recursive call

let loEnd = mid + 1; // first index such that sum(arr[i:loEnd + 1]) is in range
let hiEnd = mid + 1; // last index such that sum(arr[i:hiEnd]) is in range

// count number of ranges with i in left half and loEnd/hiEnd in right half
for (let i = lo; i <= mid; i++) {
while (loEnd <= hi && arr[loEnd] - arr[i] < lower) { // while sum(arr[i:loEnd + 1]) < lower
loEnd++;
}
hiEnd = Math.max(hiEnd, loEnd); // hiEnd >= loEnd always
while (hiEnd <= hi && arr[hiEnd] - arr[i] <= upper) { // while sum(arr[i:hiEnd + 1]) <= upper
hiEnd++;
}
count += hiEnd - loEnd; // sum(arr[i:loEnd + 1]) to sum(arr[i:hiEnd]) all in range
}

merge(lo, mid, hi, arr);
return count;
}

return mergeSort(0, prefixSum.length - 1, prefixSum);
};

We will implement a divide and conquer solution. The solution is difficult to understand. We use a modified merge sort algorithm to sort prefixSum (named arr in the mergeSort and merge method) where prefixSum[i] is the sum(nums[0:i]), while counting the number of ranges at the same time. In the sort step, we count all the number of ranges such that:

  • The start index is in [lo, mid].
  • The end index is in (mid, hi].

We do not need to worry about the other cases, since our recursive call should have already accounted for those. To actually find the counts of the above scenario, we iterate using i from lo to mid, where for each i we find:

  • Index loEnd such that sum(arr[i:loEnd + 1]) is in range.
  • Index hiEnd such that sum(arr[i:hiEnd]) is in range.

Thus, sum(arr[i:loEnd + 1]) to sum(arr[i:hiEnd]) must all be in range, and so count is incremented by hiEnd - loEnd. This works since arr is monotonically increasing in both the left and right half (ie. arr[lo:mid + 1] and arr[mid + 1:hi + 1]).

More specifically, this property ensures that the elements after loEnd and hiEnd - 1 cannot be less than loEnd and hiEnd - 1 respectively; which guarantees all possible sums outside of [loEnd, hiEnd) as the ending index (provided that i is the starting index) are not in range.

Furthermore, this property also means as i increases, the new values of loEnd and hiEnd must either stay the same or also increase (ie. we don't need to reset the values of loEnd and hiEnd back to mid + 1 after each iteration of i).

Note that arr can be sorted without impacting count. This is because although it changes the order of values in arr, there is still enough information to determine the existence of sums that are in the range [lower, upper]. For example, if we know the left half of arr contains k1 and the right half contains k2, then there must be some index i, j such that arr[i:j + 1] = k2 - k1 (where i is from the left half and j is from the right half). Even if k1 and k2 are shuffled around in their respective halves, such i, j will still exist (and so we can still account for its existence). Of course, this is no longer the case after the merge method is ran, which is why we compute count beforehand.