问题描述

给你一个整数数组 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