树状数组-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

原创文章,作者:3628473679,如若转载,请注明出处:https://blog.ytso.com/271672.html

(0)
上一篇 2022年7月6日
下一篇 2022年7月6日

相关推荐

发表回复

登录后才能评论