【Leetcode】327. Count of Range Sum 327. 区间和的个数

最新推荐文章于 2023-05-19 22:19:53 发布

MYSDB

最新推荐文章于 2023-05-19 22:19:53 发布

阅读量873

点赞数 1

分类专栏： Leetcode

本文链接：https://blog.csdn.net/lemonmillie/article/details/89055097

版权

Leetcode 专栏收录该内容

129 篇文章 1 订阅

订阅专栏

解法

naive的方法是 $O(n^2)$ 的
要找更高效的方法，都离不开先计算出一个前缀和数组p，然后对于每个元素p[i]，要计算i以后有多少个元素在[p[i]+lower,p[i]+upper]范围内

解法一：归并排序

在merge的过程中会得到两个有序的数组L，R
对于L里的每个元素a，需要在R里找到[a+lower,a+upper]范围里的元素的个数，由于R是有序的，所以用尺取法就可以在 $O (n)$ 时间内完成，所以最后复杂度为 $O (n l o g n)$

class Solution(object):
    def countRangeSum(self, nums, lower, upper):
        """
        :type nums: List[int]
        :type lower: int
        :type upper: int
        :rtype: int
        """
        from bisect import bisect_left,bisect_right,insort
        if not nums or upper<lower:
            return 0
        p = [0]
        for a in nums:
            p.append(p[-1]+a)
        n = len(p)
        self.ans = 0
        def mergesort(A):
            if len(A)<=1:
                return A
            mid = len(A)>>1
            B,C = mergesort(A[:mid]),mergesort(A[mid:])
            n = len(C)
            l = r = 0
            for b in B:
                lo,hi = b+lower,b+upper
                while l<n and C[l]<lo:
                    l += 1
                while r<n and C[r]<=hi:
                    r += 1
                self.ans += r-l
            res = []
            i = j = 0
            while i<mid:
                while j<n and C[j]<B[i]:
                    res.append(C[j])
                    j += 1
                res.append(B[i])
                i += 1
            while j<n:
                res.append(C[j])
                j += 1
            return res
        mergesort(p)
        return self.ans

解法二：TreeMap

倒序遍历前缀和数组，当遍历到p[i]的时候，它往后的元素都已经遍历过，它们都插入了一个有序的TreeMap，查找两个边界需要 $O (l o g n)$ ，然后再插入p[i]需要 $O (l o g n)$ ，问复杂度为 $O (n l o g n)$
由于python没有TreeMap的结构，所以用bisect来代替

class Solution(object):
    def countRangeSum(self, nums, lower, upper):
        """
        :type nums: List[int]
        :type lower: int
        :type upper: int
        :rtype: int
        """
        from bisect import bisect_left,bisect_right,insort
        if not nums or upper<lower:
            return 0
        p = [0]
        for a in nums:
            p.append(p[-1]+a)
        walked = []
        ans = 0
        for a in p[::-1]:
            l,r = a+lower,a+upper
            i,j = bisect_left(walked,l),bisect_right(walked,r)
            ans += j-i
            insort(walked,a)
            # print walked
        return ans

解法三：离散化树状数组

这种统计在某个区间内的数字的个数，也可以用树状数组来解决
由于数字的范围可能很大，所以要离散化，涉及到的需要被离散化的值有p[i],p[i]+lower,p[i]+upper。
树状数组的查询是 $O (l o g n)$ 的，所以最后复杂度为 $O (n l o g n)$

class Solution(object):
    def countRangeSum(self, nums, lower, upper):
        """
        :type nums: List[int]
        :type lower: int
        :type upper: int
        :rtype: int
        """
        if not nums or upper<lower:
            return 0
        p = [0]
        for a in nums:
            p.append(p[-1]+a)
        value = set()
        for a in p:
            value |= {a,a+lower,a+upper}
        value = sorted(value)
        idx={}
        for i,a in enumerate(value):
            idx[a]=i
        n = len(value)
        c = [0]*n
        def lowbit(k):
            return k&-k
        def add(k):
            k +=1
            while k<=n:
                c[k-1]+=1
                k+=lowbit(k)
        def tsum(k):
            res = 0
            k+=1
            while k>=1:
                res += c[k-1]
                k -= lowbit(k)
            return res
        ans = 0
        for a in p[::-1]:
            l,r = idx[a+lower],idx[a+upper]
            ans += tsum(r)
            if l>0:
                ans -= tsum(l-1)
            add(idx[a])
        return ans