lc 327. count of range sum

Description

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

A very naive solution is to use O(n^2) time complexity after getting the accumulation of A[:i] array.

We can use a merge sort solution(divide and conqure) to deal with it. We assume that sums[lo:mi] and sums[mi:lo] are sorted, so we can traverse i from lo to mi and for every i, we should find how many j in mi to hi that can fullfill sums[j] - sums[i] is in the [lower, upper] range.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class :
def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
def countWhileMergeSort(lo, hi):
nonlocal sums
if hi - lo <= 1: return 0
mi = (hi - lo) // 2 + lo
cache = [0] * (hi - lo)
j, k, t, r = mi, mi, mi, 0
res = countWhileMergeSort(lo, mi) + countWhileMergeSort(mi, hi)
for i in range(lo, mi):
while j < hi and sums[j] - sums[i] < lower: j += 1
while k < hi and sums[k] - sums[i] <= upper: k += 1
while t < hi and sums[t] < sums[i]:
cache[r] = sums[t]
r, t = r + 1, t + 1
cache[r] = sums[i]
r += 1
res += k - j
for i in range(lo, t):
sums[i] = cache[i - lo]
return res
n = len(nums)
sums = [0] * (n + 1)
for i in range(n):
sums[i + 1] = sums[i] + nums[i]
return countWhileMergeSort(0, n + 1)