327. Count of Range Sum
中午请人吃饭,结果吃多了,好困,有点坐不动了。这题有segment tree其实不是那么好做,答案的解释一大堆,决定放到明天去理解。今天的任务是做到hard70,大概是可以完成的。
class Solution(object):
def merge(self, arr1, arr2):
r, i, j = [], 0, 0
while i < len(arr1) and j < len(arr2):
if arr1[i] < arr2[j]:
r.append(arr1[i])
i += 1
else:
r.append(arr2[j])
j += 1
r += arr1[i:] + arr2[j:]
return r
def count(self, prefix, start, end, lower, upper):
if start >= end:
return 0
mid = start + (end - start + 1 >> 1)
count = sum([
self.count(prefix, s, e, lower, upper)
for s, e in ((start, mid - 1), (mid, end))
])
l, r = mid, mid
for i in xrange(start, mid):
while l <= end and prefix[l] - prefix[i] < lower:
l += 1
while r <= end and prefix[r] - prefix[i] <= upper:
r += 1
count += r - l
prefix[start:end + 1] = self.merge(
prefix[start:mid], prefix[mid:end + 1])
return count
def countRangeSum(self, nums, lower, upper):
n = len(nums)
prefix = [0] * (n + 1)
for i, v in enumerate(nums, 1):
prefix[i] = prefix[i - 1] + v
return self.count(prefix, 0, n, lower, upper)