题目描述:
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:
Given nums = [-2, 5, -1]
, lower = -2
, upper = 2
,
Return 3
.
The three ranges are : [0, 0]
, [2, 2]
, [0, 2]
and their respective sums are: -2, -1, 2
.
题目大意:
给定一个整数数组nums,返回其所有落在[low, upper]范围内(包含边界)的区间和的数目。
区间和S(i, j)的定义为所有下标为i到j之间(i ≤ j)的元素的和,包含边界。
注意:
朴素的O(n^2)解法很容易想到。你的算法必须优于朴素解法。
测试用例如题目描述。
解题思路:
本题可以类比题目Count of Smaller Numbers After Self的解法
朴素解法O(n^2):
预处理前n项和数组sums,其中sums[i] = ∑(nums, 0, i)
则S(i, j) = sums[j] - sums[i - 1]
两重循环枚举i, j即可。
O(n * logn)解法:
解法I 树状数组(Fenwick Tree):
1. 预处理前n项和数组sums 2. 将sums数组离散化(排序+去重)得到数组osums 3. 遍历sums,记sumi = sums[i] 用二分查找得到[sumi - upper, sumi - lower]的离散化下标[left, right] 用树状数组统计范围[left, right]内的元素个数,并累加至最终结果ans 若lower <= sumi <= upper,额外地令ans+1 将sumi的离散化下标记入树状数组
Python代码:
class Solution(object):
def countRangeSum(self, nums, lower, upper):
"""
:type nums: List[int]
:type lower: int
:type upper: int
:rtype: int
"""
sums = nums[:]
for x in range(1, len(sums)):
sums[x] += sums[x - 1]
osums = sorted(set(sums))
ft = FenwickTree(len(osums))
ans = 0
for sumi in sums:
left = bisect.bisect_left(osums, sumi - upper)
right = bisect.bisect_right(osums, sumi - lower)
ans += ft.sum(right) - ft.sum(left) + (lower <= sumi <= upper)
ft.add(bisect.bisect_right(osums, sumi), 1)
return ans
class FenwickTree(object):
def __init__(self, n):
self.n = n
self.sums = [0] * (n + 1)
def add(self, x, val):
while x <= self.n:
self.sums[x] += val
x += self.lowbit(x)
def lowbit(self, x):
return x & -x
def sum(self, x):
res = 0
while x > 0:
res += self.sums[x]
x -= self.lowbit(x)
return res
解法II 归并排序(Merge Sort):
参考思路:https://leetcode.com/discuss/79083/share-my-solution
Python代码:
class Solution(object):
def countRangeSum(self, nums, lower, upper):
"""
:type nums: List[int]
:type lower: int
:type upper: int
:rtype: int
"""
size = len(nums)
sums = [0] * (size + 1)
for x in range(1, size + 1):
sums[x] += nums[x - 1] + sums[x - 1]
INF = max(sums)
def mergeSort(lo, hi):
if lo == hi: return 0
mi = (lo + hi) / 2
cnt = mergeSort(lo, mi) + mergeSort(mi + 1, hi)
x = y = lo
for i in range(mi + 1, hi + 1):
while x <= mi and sums[i] - sums[x] >= lower:
x += 1
while y <= mi and sums[i] - sums[y] > upper:
y += 1
cnt += x - y
part = sums[lo : hi + 1]
l, h = lo, mi + 1
for i in range(lo, hi + 1):
x = part[l - lo] if l <= mi else INF
y = part[h - lo] if h <= hi else INF
if x < y: l += 1
else: h += 1
sums[i] = min(x, y)
return cnt
return mergeSort(0, size)