327-count-of-range-sum.md

327. 区间和的个数 - hard

给定一个整数数组 nums 。区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。

请你以下标 i （0 <= i <= nums.length ）为起点，元素个数逐次递增，计算子数组内的元素和。

当元素和落在范围 [lower, upper] （包含 lower 和 upper）之内时，记录子数组当前最末元素下标 j ，记作 有效 区间和 S(i, j) 。

求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 有效 区间和的个数。

 

注意：
最直观的算法复杂度是 O(n²) ，请在此基础上优化你的算法。

 

示例：

输入：nums = [-2,5,-1], lower = -2, upper = 2,
输出：3 
解释：
下标 i = 0 时，子数组 [-2]、[-2,5]、[-2,5,-1]，对应元素和分别为 -2、3、2 ；其中 -2 和 2 落在范围 [lower = -2, upper = 2] 之间，因此记录有效区间和 S(0,0)，S(0,2) 。
下标 i = 1 时，子数组 [5]、[5,-1] ，元素和 5、4 ；没有满足题意的有效区间和。
下标 i = 2 时，子数组 [-1] ，元素和 -1 ；记录有效区间和 S(2,2) 。
故，共有 3 个有效区间和。

 

提示：

• 0 <= nums.length <= 10^4

Solutions

1. 前缀和数组 + 暴力

假设前缀和数组为 prefix_sum, 那么一定有 j >= i 且 prefix_sum[j] - prefix_sum[i] ∈ [lower, height]。

先构建前缀和数组，再暴力计算满足 i, j 对的个数，就是结果。

空间复杂度为 O(n) 时间复杂度为 O(n^2)

会超时。

class Solution:
    def countRangeSum(self, nums: List[int], lower: int, upper: int) -> int:
        prefix_sum = [0 for _ in range(len(nums) + 1)]
        for i, v in enumerate(nums):
            prefix_sum[i + 1] = prefix_sum[i] + v

        cnt = 0
        for i in range(len(prefix_sum)):
            j = i + 1
            while j < len(prefix_sum):
                v = prefix_sum[j] - prefix_sum[i]
                if lower <= v <= upper:
                    cnt += 1
                j += 1

        return cnt

2. 前缀和数组 + 归并排序

里面关键的思想是，利用归并排序分治过程。来往后面比较。分治过程中，会将数组分为两部分。左边所有元素都在右边元素之前。将左边每一个元素与右边元素做差值比较就行。又因为左右两组数据都是有序的。所以迭代左边数组时，left 和 right 指针只需要往从左往右走一遍就行。

左边和左边内部，会在子的归并排序中分为两部分，同上逻辑做判断。

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

        def merge_count(arr: List[int]) -> int:
            n = len(arr)
            if n <= 1:
                return 0

            n1 = arr[:int(n/2)]
            n2 = arr[int(n/2):]
            cnt = merge_count(n1) + merge_count(n2)

            l, r = 0, 0
            for v in n1:
                while l < len(n2) and n2[l] - v < lower:
                    l += 1
                while r < len(n2) and n2[r] - v <= upper:
                    r += 1
                cnt += (r - l)
            
            p1, p2 = 0, 0
            for i, _ in enumerate(arr):
                if p1 < len(n1) and (p2 == len(n2) or n1[p1] <= n2[p2]):
                    arr[i] = n1[p1]
                    p1 += 1
                else:
                    arr[i] = n2[p2]
                    p2 += 1
            
            return cnt
        
        prefix_sum = [0 for _ in range(len(nums) + 1)] 
        for i, v in enumerate(nums):
            prefix_sum[i + 1] = prefix_sum[i] + v
        
        return merge_count(prefix_sum)