Python | Leetcode Python题解之第327题区间和的个数

题目：

题解：

class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        pre = list(accumulate(nums, initial=0))
        nums = sorted(pre)
        mx = len(nums)
        b = BIT(mx + 1)
        ans = 0
        # 统计[x-upper,x-lower]的个数
        for i, x in enumerate(pre):
            j = bisect_left(nums, x) + 1
            r = bisect_right(nums, x - lower)    # <= x - lower , 注意原先应该-1，但是下标从1开始再+1
            l = bisect_left(nums, x - upper)     # < x - lower , 注意原先应该-1，但是下标从1开始再+1
            ans += b.query(r) - b.query(l)
            b.add(j, 1)
        return ans

class BIT:
    def __init__(self, n: int):
        self.tree = [0] * n  # 树状数组
        self.original = [0] * n  # 原数组

    def update(self, i: int, val: int) -> None:
        self.original[i] = max(self.original[i], val)
        while i < len(self.tree):
            self.tree[i] = max(self.tree[i], val)
            i += i & -i

    def query_max(self, L: int, R: int) -> int:
        mx = 0
        while R >= L:
            r = R & (R - 1)
            # 查询先进行比较，看下一个r在不在查询范围内
            if r >= L:
                # 在查询范围内，直接从树状数组拿值比较
                mx = max(mx, self.tree[R])
                R = r
            else:
                # 只走一步，从原数组拿值比较
                mx = max(mx, self.original[R])
                R -= 1
        return mx

    # 统计 <= R 的元素个数
    def query(self, R: int) -> int:
        res = 0
        while R > 0:
            res += self.tree[R]
            R &= (R - 1)
        return res

    def add(self, i: int, val: int) -> None:
        while i < len(self.tree):
            self.tree[i] += val
            i += i & -i