【Leetcode】327.区间和的个数

327 区间和的个数

题目

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

示例:
输入: nums = [-2,5,-1], lower = -2, upper = 2,
输出: 3 
解释: 3个区间分别是: [0,0], [2,2], [0,2]，它们表示的和分别为: -2, -1, 2。

即求满足累加和在lower和upper之间的子数组的个数

方法1（前缀和）

首先生成前缀和数组preSum，其中preSum[i]表示nums中i位置之前所有元素之和

1
2
3
4
5
long[] preSum = new long[nums.length + 1];
preSum[0] = 0;
for(int i = 1; i <= nums.length; i++){
    preSum[i] = preSum[i - 1] + nums[i - 1];
}

可以看出，子数组nums[i….j-1]的累加和可以由前缀和数组快速地得到，即preSum[j] - preSum[i]。我们可以遍历出所有可能的子数组，再判断其累加和是否在lower和upper之间即可，但这种方法的时间复杂度为O(n^2)

方法2（前缀和+归并排序）

我们得到了前缀和数组preSum后，如果preSum[j] - preSum[i]大于lower小于upper，那么就说明nums[i….j-1]的累加和大于lower小于upper。

如果考虑把前缀和数组分为两段，考虑前一段和后一段都有序、两段之间无序的情况。如果我们先固定前一段的一个元素i，假设后一段的索引为j，我们如何得到一个j的取值区间，让区间内preSum[j] - preSum[i]都大于lower小于upper呢？

  • 先定义两个变量L和R，分别代表这个区间的左右边界。并将L和R都初始化指向后一段的第一个元素
  • 如果这时preSum[j] - preSum[i]还不大于lower，则说明还没找到这样的区间，区间的左边界也就不在此，L向右移。直到遇到preSum[j] - preSum[i]大于lower的j位置，说明我们已经进入了这个区间，此时的L即为区间的左边界。
  • 当前preSum[L] - preSum[i]大于lower，又因为后半段有序，因此L位置之后的所有j都满足preSum[j] - preSum[i]大于lower
  • 接下来需要寻找满足preSum[j] - preSum[i]小于upper的j的右边界，步骤同上，我们用R记录下来这个右边界。至此为止，L和R间的所有j位置都满足preSum[j] - preSum[i]大于lower小于upper。

因此我们知道，以i位置开头的，满足累加和大于lower小于upper的子数组一共有R-L个。而又因为左半段有序，因此preSum[i+1]大于preSum[i],若要寻找以i+1开头的子数组，则之前的L和R必然只能向右移动。同理，我们可以找到以左半段每个元素开头的，满足累加和大于lower小于upper的子数组个数。加起来就是当i在左半段、j在右半段的情况下，满足preSum[j] - preSum[i]大于lower小于upper的组合个数。

而结果需要包含三种情况：

  • i和j都在左半段
  • i和j都在右半段
  • i在左半段、j在右半段

我们刚才所求的是第三种，第一种和第二种情况我们可以将左半段再分两段，右半段再分两段，递归地求得。（归并排序的思想）

时间复杂度和归并排序相同，为O(nlogn)

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class Solution {
    private int lower;
    private int upper;
    public int countRangeSum(int[] nums, int lower, int upper) {
        this.lower = lower;
        this.upper = upper;
        if (nums == null || nums.length == 0) {
            return 0;
        }
        //求前缀和数组
        long[] preSum = new long[nums.length + 1];
        preSum[0] = 0;
        for(int i = 1; i <= nums.length; i++){
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
        return mergeSort(preSum, 0, preSum.length - 1);
    }

    private int mergeSort(long[] preSum, int left, int right) {
        if (left >= right) {
            return 0;
        }
        int mid = left + ((right - left) >> 1);
        //满足累加和在lower和upper之间的子数组的个数分为三部分：
        //左侧自己merge时找到的子数组个数（i、j都在左侧）
        //右侧自己merge时找到的子数组个数（i、j都在右侧）
        //左侧和右侧merge时找到的子数组个数（i在左侧、j在右侧）
        return mergeSort(preSum, left, mid) + mergeSort(preSum, mid + 1, right) + merge(preSum, left, mid, right);
    }

    private int merge(long[] preSum, int left, int mid, int right) {
        // 统计i在左半段、j在右半段的情况下满足条件的子数组个数
        // 左半段为[left, mid]，右半段为[mid + 1, right]
        // j的左右边界分别为L和R
        int res = 0;
        int i = left;
        int L = mid + 1, R = mid + 1;
        while (i <= mid) {
            while (L <= right && preSum[L] - preSum[i] < lower) 
                L++;
            while (R <= right && preSum[R] - preSum[i] <= upper) 
                R++;
            res += R - L; 
            i++;
        }

        //进行归并排序的标准merge操作
        int[] help = new int[right - left + 1];
        int p1 = left, p2 = mid + 1;
        int idx = 0;

        while (p1 <= mid && p2 <= right) {
            help[idx++] = (preSum[p1] <= preSum[p2]) ? (int) preSum[p1++] : (int) preSum[p2++];
        }
        while (p1 <= mid) {
            help[idx++] = (int) preSum[p1++];
        }
        while (p2 <= right) {
            help[idx++] = (int) preSum[p2++];
        }
        for (int j = 0; j < help.length; j++) {
            preSum[left + j] = help[j];
        }
        return res;
    }
}