题目:
题解:
class Solution:
def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
pre = list(accumulate(nums, initial=0))
nums = sorted(pre)
mx = len(nums)
b = BIT(mx + 1)
ans = 0
# 统计[x-upper,x-lower]的个数
for i, x in enumerate(pre):
j = bisect_left(nums, x) + 1
r = bisect_right(nums, x - lower) # <= x - lower , 注意原先应该-1,但是下标从1开始再+1
l = bisect_left(nums, x - upper) # < x - lower , 注意原先应该-1,但是下标从1开始再+1
ans += b.query(r) - b.query(l)
b.add(j, 1)
return ans
class BIT:
def __init__(self, n: int):
self.tree = [0] * n # 树状数组
self.original = [0] * n # 原数组
def update(self, i: int, val: int) -> None:
self.original[i] = max(self.original[i], val)
while i < len(self.tree):
self.tree[i] = max(self.tree[i], val)
i += i & -i
def query_max(self, L: int, R: int) -> int:
mx = 0
while R >= L:
r = R & (R - 1)
# 查询先进行比较,看下一个r在不在查询范围内
if r >= L:
# 在查询范围内,直接从树状数组拿值比较
mx = max(mx, self.tree[R])
R = r
else:
# 只走一步,从原数组拿值比较
mx = max(mx, self.original[R])
R -= 1
return mx
# 统计 <= R 的元素个数
def query(self, R: int) -> int:
res = 0
while R > 0:
res += self.tree[R]
R &= (R - 1)
return res
def add(self, i: int, val: int) -> None:
while i < len(self.tree):
self.tree[i] += val
i += i & -i