树状数组-327. 区间和的个数
问题描述
给你一个整数数组 nums 以及两个整数 lower 和 upper 。求数组中,值位于范围 [lower, upper] (包含 lower 和 upper)之内的 区间和的个数 。
区间和 S(i, j) 表示在 nums 中,位置从 i 到 j 的元素之和,包含 i 和 j (i ≤ j)。
示例 1:
输入:nums = [-2,5,-1], lower = -2, upper = 2
输出:3
解释:存在三个区间:[0,0]、[2,2] 和 [0,2] ,对应的区间和分别是:-2 、-1 、2 。
示例 2:
输入:nums = [0], lower = 0, upper = 0
输出:1
提示:
1 <= nums.length <= 105
-231 <= nums[i] <= 231 - 1
-105 <= lower <= upper <= 105
题目数据保证答案是一个 32 位 的整数
问题求解
(j, i]: presum[i] - presum[j]
对于presum[i],求解满足: lower <= presum[i] - presum[j] <= upper and j < i的j的个数
等价于 presum[i] - uppper <= presum[j] <= presum[i] - lower and j < i的j的个数
因此可以使用排序离散化,计频数组求和的方式来解决。
需要主要的有两点:1)离散化的时候需要将所有涉及的数字进行离散化;2)树状数组的更新要在求解后。
from sortedcontainers import SortedSet
class Solution:
def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
n = len(nums)
presum = [0] * n
presum[0] = nums[0]
for i in range(1, n):
presum[i] = presum[i - 1] + nums[i]
# (j, i]: presum[i] - presum[j]
# 对于presum[i],求解满足: lower <= presum[i] - presum[j] <= upper and j < i的j的个数
# 等价于 presum[i] - uppper <= presum[j] <= presum[i] - lower and j < i的j的个数
# 离散化
sort = SortedSet()
for num in presum:
sort.add(num)
sort.add(num - lower)
sort.add(num - upper)
record = {}
rank = 1
for num in sort:
record[num] = rank
rank += 1
bit = [0] * (len(record) + 1)
def update(idx, delta):
while idx < len(bit):
bit[idx] += delta
idx += (idx & -idx)
def query(idx):
res = 0
while idx:
res += bit[idx]
idx -= (idx & -idx)
return res
res = 0
for i in range(n):
if lower <= presum[i] <= upper:
res += 1
# 由于j < i,因此要讲update放在后面,不然就将本轮的数加入了j的计算当中
r, l = record[presum[i] - lower], record[presum[i] - upper]
res += query(r) - query(l - 1)
update(record[presum[i]], 1)
return res