标签: 树状数组 线段树 数组 二分查找 分治 有序集合 归并排序
难度: Hard
给你一个整数数组 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运行时间: 704 ms
内存: 32.1 MB
class Solution:
def countRangeSum(self, nums, lower, upper):
arr_sum = [0] * len(nums) # 前缀和数组
arr_sum[0] = nums[0]
for i in range(1, len(nums)):
arr_sum[i] = arr_sum[i - 1] + nums[i]
return self.process5(arr_sum, 0, len(arr_sum) - 1, lower, upper)
def process5(self, arr, left, right, lower, upper):
if left == right:
return 1 if lower <= arr[left] <= upper else 0
mid = (left + right) // 2
return sum([self.process5(arr, left, mid, lower, upper), self.process5(arr, mid + 1, right, lower, upper),
self.merge5(arr, left, mid, right, lower, upper)])
def merge5(self, arr, left, mid, right, lower, upper):
ans = 0
help_arr = [0] * (right - left + 1) # 辅助数组，就两个数组合并，最后再赋值给arr[left:right+1]
i = 0
p1 = left
p2 = mid + 1
# [window_l, window_r) 是囊括的答案数
window_l = left
window_r = left
for v in arr[mid + 1: right + 1]:
value_min = v - upper
value_max = v - lower
while window_l <= mid and arr[window_l] < value_min: # l 最大能到达 m+1
window_l += 1
while window_r <= mid and arr[window_r] <= value_max: # r 最大能到达 m+1
window_r += 1
ans += window_r - window_l
while (p1 <= mid) and (p2 <= right):
if arr[p1] <= arr[p2]: #
help_arr[i] = arr[p1]
p1 += 1
else: # 等于的时候优先让右边进辅助数组
help_arr[i] = arr[p2]
p2 += 1
i += 1
help_arr[i:] = arr[p1:mid + 1] if p1 <= mid else arr[p2:right + 1] # 等价于上面两个while
arr[left: right + 1] = help_arr
return ans
该题解采用分治法与归并排序的思想来解决问题。首先，我们计算了数组的前缀和，然后通过递归分治方法处理每个子数组，并在合并的过程中统计符合条件的区间和的个数。具体步骤如下： 1. 计算数组的前缀和，以便快速计算任意子数组的和。 2. 使用递归分治的方法来处理和合并子数组。对于每个子问题，如果它只包含一个元素，直接判断这个单独元素形成的区间和是否在指定范围内。 3. 在合并两个已经排序的子数组的过程中，使用双指针技术统计当前右侧元素与左侧元素组合形成的区间和落在[lower, upper]范围内的个数。 4. 合并完成后，还需要将左右两部分归并排序，以便后续操作可以继续利用已排序的性质加速区间和的计算。
时间复杂度: O(n log n)
空间复杂度: O(n)
class Solution:
def countRangeSum(self, nums, lower, upper):
arr_sum = [0] * len(nums) # 计算前缀和数组
arr_sum[0] = nums[0]
for i in range(1, len(nums)):
arr_sum[i] = arr_sum[i - 1] + nums[i]
return self.process5(arr_sum, 0, len(arr_sum) - 1, lower, upper)
def process5(self, arr, left, right, lower, upper):
if left == right:
return 1 if lower <= arr[left] <= upper else 0
mid = (left + right) // 2
return sum([self.process5(arr, left, mid, lower, upper), self.process5(arr, mid + 1, right, lower, upper),
self.merge5(arr, left, mid, right, lower, upper)])
def merge5(self, arr, left, mid, right, lower, upper):
ans = 0
help_arr = [0] * (right - left + 1) # 辅助数组，用于合并两个子数组
i = 0
p1 = left
p2 = mid + 1
window_l = left
window_r = left
for v in arr[mid + 1: right + 1]:
value_min = v - upper
value_max = v - lower
while window_l <= mid and arr[window_l] < value_min:
window_l += 1
while window_r <= mid and arr[window_r] <= value_max:
window_r += 1
ans += window_r - window_l
while (p1 <= mid) and (p2 <= right):
if arr[p1] <= arr[p2]:
help_arr[i] = arr[p1]
p1 += 1
else:
help_arr[i] = arr[p2]
p2 += 1
i += 1
help_arr[i:] = arr[p1:mid + 1] if p1 <= mid else arr[p2:right + 1] # 处理剩余元素
arr[left: right + 1] = help_arr # 将排序后的数组写回原数组
return ans
使用前缀和可以将计算任意子数组的和的时间复杂度从线性时间降低到常数时间。如果直接遍历所有区间组合来计算区间和，则对于长度为 n 的数组，可能的区间组合数为 n*(n+1)/2，这将导致时间复杂度为 O(n^2)，这在 n 较大时效率非常低。通过计算前缀和数组，我们只需 O(n) 时间就可以得到整个数组的前缀和，之后计算任何一个区间和都可以在 O(1) 时间内完成，从而大大提高效率。
在这个问题中，分治法通过递归地将数组分成更小的子数组，直到子数组只包含一个元素，然后在合并过程中计算跨越两个子数组的区间和是否在给定的范围内。具体来说，我们首先计算左半部分和右半部分内部各自满足条件的区间和的数量，然后计算跨越左右子数组的区间和的数量。在合并步骤中，我们使用双指针技术：一个指针遍历右子数组，另外两个指针在左子数组中移动，定位符合条件的区间和范围。这种方法利用了区间和的性质，即如果我们知道右子数组的某个前缀和，可以快速通过左子数组的前缀和确定哪些和落在指定范围内。这种分治加合并的策略有效地减少了重复计算，从而提高了算法的效率。
在合并过程中，使用双指针来统计符合条件的区间和的数量主要涉及遍历右子数组的每个元素，并对每个元素使用两个指针在左子数组中寻找符合条件的前缀和范围。这个过程的时间复杂度是 O(n)，因为虽然对每个右子数组元素都可能进行一系列比较，但是两个指针（window_l 和 window_r）都是单向移动，即每个指针在整个合并过程中最多移动 n 步。因此，这一步的总体操作数是线性的，即 O(n)。这比直接计算每个可能的区间和要高效得多。