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 thatsum(arr[i:loEnd + 1])
is in range. - Index
hiEnd
such thatsum(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.