↓Skip to main content
  1. Posts/

LeetCode-327 区间和的个数

·1 min·
Table of Contents

LeetCode-327 区间和的个数 #

Solution 1 #

对于区间和, 转化成前缀和之差 (需要注意在前缀和数组最前面保留一个 $0$, 不遗漏从头开始的完整区间) . 对于前缀和 $presum[i]$ , 需要对满足 $presum[i] + lower \leq presum[j] \leq presum[i] + upper$ 的 $j$ 进行计数. 考虑有两个有序数组, 对于左侧的每个 $i$, 不断维护右侧 $l, r$ 双指针即可. 我们可以使用归并排序来保证这一点. 代码如下:

class Solution:
    ans = 0
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        self.ans = 0

        presum = [0 for _ in range(len(nums) + 1)]

        for i, x in enumerate(nums):
            presum[i + 1] = x + presum[i]

        def merge_sort(a):
            n = len(a)
            if n <= 1:
                return a
            n1 = n // 2
            n2 = n - n1
            a1, a2 = a[:n1], a[n1:]
            a1, a2 = merge_sort(a1), merge_sort(a2)

            l, r = 0, 0
            for i in range(n1):
                while l < n2 and a2[l] - a1[i] < lower:
                    l += 1

                while r < n2 and a2[r] - a1[i] <= upper:
                    r += 1

                # [l, r) 里面是满足计数条件的
                self.ans += r - l

            res = []
            p1, p2 = 0, 0
            while p1 < len(a1) and p2 < len(a2):
                if a1[p1] <= a2[p2]: 
                    res.append(a1[p1])
                    p1 += 1
                else:
                    res.append(a2[p2])
                    p2 += 1

            while p1 < len(a1): 
                res.append(a1[p1])
                p1 += 1
                
            while p2 < len(a2):
                res.append(a2[p2])
                p2 += 1

            return res

        merge_sort(presum)
        return self.ans
↑