首页手记 327. 区间和的个数

327. 区间和的个数

标签：

Python 算法

@TOC

327. 区间和的个数

一、题目

1. 题目描述

给你一个整数数组 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 <= 10⁵
-2³¹ <= nums[i] <= 2³¹ - 1
-10⁵ <= lower <= upper <= 10⁵
题目数据保证答案是一个 32 位 的整数

2. 原题链接

327. 区间和的个数:

二、解题报告

1. 思路分析

题意表述不清，这里重新说明：找到数组中所有区间（子数组），对每个区间求和。问有多少个符合条件的区间(求和在lower和upper之间)。

看到是区间和问题，先想到前缀和。预处理出一个前缀和数组

s = list(accumulate(nums,initial=0))

区间[i,j]的和为s[j]-s[i]
对每个j，问题转化为它左边的i有多少个满足条件的：lower<=s[j]-s[i]<=upper
上述式子转化为，s[j]-upper<=s[i]<=s[j]-lower
那么只需要遍历j(让j做常量),问左边有多少个s[i]符合条件s[j]-upper<=s[i]<=s[j]-lower
由于j遍历时是一直增加的，相当于一直有update和query，那么有没有一种数据结构可以实现低复杂度的查询和更新呢。
我们知道树状数组和线段树的查询和更新都是O(nlog₂n)
参考步骤四，我们的数据结构要维护的是数字的数量，没插入一个，这个数数量+1，那么线段的坐标代表的是数字本身。
由于数据范围大，需要对数据进行离散化，包括s[j]-upper，s[i]，s[j]-lower

2. 复杂度分析

最坏时间复杂度O(nlog₂n)

3.代码实现

线段树：

class IntervalTree:
    def __init__(self, size):
        self.size = size
        self.interval_tree = [0 for _ in range(size*4)]

    def insert(self,p,l,r,index):
        interval_tree = self.interval_tree        
        if l == r:
            interval_tree[p] += 1
            return
        mid = (l+r)//2
        if index <= mid:
            self.insert(p*2,l,mid,index)
        else:
            self.insert(p*2+1,mid+1,r,index)
        interval_tree[p] = interval_tree[p*2]+interval_tree[p*2+1]       
    
    def query(self,p,l,r,x,y):
        if x<=l and r<=y:
            return self.interval_tree[p]
        mid = (l+r)//2
        s = 0
        if x <= mid:
            s += self.query(p*2,l,mid,x,y)
        if mid < y:
            s += self.query(p*2+1,mid+1,r,x,y)
        return s

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 = IntervalTree(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.query(1,1,tree_size, x+1,y+1)
            # print(x,y,j,c)
            cnt += c
            tree.insert(1,1,tree_size,j+1)
            
        return cnt

树状数组：

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