[英雄星球七月集训LeetCode解题日报] 第25日 树状数组

日报

• 今天的题很多种解法。还是有序集合快。

题目

一、 327. 区间和的个数

链接: 1327. 区间和的个数

1. 题目描述

2. 思路分析

• 多种解法，本质是转化为前缀和后统计s[j]-upper<=s[i]<=s[j]-lower,统计这些i的数量。
• 因此可以用PUIQ模型统计。

3. 代码实现

树状数组

class BinIndexTree:
    def __init__(self, size):
        self.size = size
        self.bin_tree = [0 for _ in range(size*4)]
    def add(self,i,v):
        while i<=self.size :
            self.bin_tree[i] += v
            i += self.lowbit(i)
    def sum(self,i):
        s = 0
        while i >= 1:
            s += self.bin_tree[i]
            i -= self.lowbit(i)
        return s

    def lowbit(self,x):
        return x&-x

    
class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:         
        s = list(accumulate(nums,initial=0))
        hashes = s + [ x-lower for x in s] + [ x-upper for x in s]
        hashes = sorted(list(set(hashes)))
        # 生成前缀和，问题转化为，对于每个j，找左边的i,判断 s[j]-upper<=s[i]<=s[j]-lower,统计这些i的数量
        # 把所有前缀和数组中的数字插入线段树，并对这些数字划分区间，线段树维护当前区间数字数量，
        # 所以需要对这些数字都散列化

        # 这里用树状数组实现上述操作
        # 树状数组也是维护每个数字出现的次数
        tree_size = len(hashes)
        tree = BinIndexTree(tree_size)
        cnt = 0
        for i in s:
            x = bisect_left(hashes,i-upper)
            y = bisect_left(hashes,i-lower)
            j = bisect_left(hashes,i)
            c = tree.sum(y+1) - tree.sum(x)
            cnt += c
            tree.add(j+1,1)

        return cnt
    

SortedList

from sortedcontainers import SortedList
class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:         
       
        # 生成前缀和，问题转化为，对于每个j，找左边的i,判断 s[j]-upper<=s[i]<=s[j]-lower,统计这些i的数量
        # 把所有前缀和数组中的数字插入有序列表，查询这个范围内的数字数量
        
        arr = SortedList()
        ans = 0
        for i in accumulate(nums,initial=0):
            ans += arr.bisect_right(i-lower)-arr.bisect_left(i-upper)
            arr.add(i)

        return ans