平衡二叉树.327. 区间和的个数

题目描述

给定一个整数数组 nums 。区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤
j)。

请你以下标 i （0 <= i <= nums.length ）为起点，元素个数逐次递增，计算子数组内的元素和。

当元素和落在范围 [lower, upper] （包含 lower 和 upper）之内时，记录子数组当前最末元素下标 j ，记作 有效
区间和 S(i, j) 。

求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 有效 区间和的个数。

来源：力扣（LeetCode） 链接：https://leetcode-cn.com/problems/count-of-range-sum
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

示例

示例：

输入：nums = [-2,5,-1], lower = -2, upper = 2, 输出：3 解释： 下标 i = 0 时，子数组
[-2]、[-2,5]、[-2,5,-1]，对应元素和分别为 -2、3、2 ；其中 -2 和 2 落在范围 [lower = -2,
upper = 2] 之间，因此记录有效区间和 S(0,0)，S(0,2) 。 下标 i = 1 时，子数组 [5]、[5,-1] ，元素和
5、4 ；没有满足题意的有效区间和。 下标 i = 2 时，子数组 [-1] ，元素和 -1 ；记录有效区间和 S(2,2) 。 故，共有
3 个有效区间和。

来源：力扣（LeetCode） 链接：https://leetcode-cn.com/problems/count-of-range-sum
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

解题思路

既然要求区间和，那么肯定需要求出前缀和。

暴力解法，就是双重循环，求出所有的 i——j 之间的和，然后进行比较。

在这里我想，任意的 i 开始，到结尾，如果是有序的，则可以通过二分进行搜索。

在向后遍历的时候，需要减去前面的和。

而既要在删除的同时又要保证有序，且时间复杂度不能高于O(n).

很容易会想到 b树。

边删除，边更新 lower 和 upper。

代码

from sortedcontainers import SortedList
class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        n = len(nums)
        presum = [0] * (n+1)
        btree = SortedList()
        for i in range(1, n+1):
            presum[i] = nums[i-1] + presum[i-1]
            btree.add(presum[i])
        res = 0
        for i in range(n):
            left = btree.bisect_left(lower)
            right = btree.bisect_right(upper)
            res += (right - left)
            btree.remove(presum[i+1])
            lower += nums[i]
            upper += nums[i]
        return res

复杂度分析

• 时间复杂度：O(nlogn)
• 空间复杂度：O(n)