327. 区间和的个数

Difficulty
Tags
URL
https://leetcode-cn.com/problems/count-of-range-sum/
Star
给你一个整数数组 nums 以及两个整数 lower 和 upper 。求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 区间和的个数 。
区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。
示例 1：
输出：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 位 的整数
通过次数28,626提交次数68,697

树状数组

思路
题目要计算区间和的个数，首先想到了前缀树组 prefix，然后再遍历前缀数组中的每一个值，假设当前遍历值为 val, 那么再去找 之前的值 在 [val-upper, val-lower] 范围内的个数。
此时是不是可以想到了树状数组，即用树状数组来维护prefix 中 val 值前面 在 [val-upper, val-lower] 范围内的个数。
需要注意的是，每次都需要查询 区间 [val-upper, val-lower] 范围内的个数，所以在树状数组中 比必须有 val - upper, val - lower for val in prefix 这些值。！！
题解
class BIT:
    def __init__(self, n):
        self.tree = [0] * (n+1)
    
    def add(self, idx, val=1):
        i = idx + 1
        while i < len(self.tree):
            self.tree[i] += val
            i += self.lowbit(i)

    def prefix(self, idx):
        i = idx + 1
        sum_ = 0
        while i > 0:
            sum_ += self.tree[i]
            i -= self.lowbit(i)

        return sum_
    def rangeSum(self, left, right):
        return self.prefix(right) - self.prefix(left-1)
    
    def lowbit(self, x):
        return x & (-x)


class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        prefix = [0] * (len(nums) + 1)

        for i in range(1, len(prefix)):
            prefix[i] = prefix[i-1] + nums[i-1]

        # 对前缀和压缩， 同时要将所有这些数字放进去
        temp = prefix[:]
        for i in prefix:
            temp.extend([i-lower, i-upper])

        uniques = sorted(list(set(temp))) 
        
        hash_map = {num:i for i, num in enumerate(uniques)}

        t = BIT(len(uniques))
        res = 0
        for i in range(len(prefix)):
            rank = hash_map[prefix[i]]

            left = hash_map[prefix[i] - upper]
            right = hash_map[prefix[i] - lower]
            res += t.rangeSum(left, right)

            t.add(rank)

        return res