使用归并排序解决「区间和的个数」问题

问题描述

给定一个整数数组 nums ，求出该数组中区间和在 [lower, upper] 之间的个数。

题目链接

https://leetcode.com/problems/count-of-range-sum/

解决思路

我们可以使用归并排序的思想来解决该问题。首先，我们定义一个辅助数组 sum ，其中 sum[i] 表示数组 nums 从下标 0 到下标 i 的元素之和。然后，我们可以通过计算 sum 数组中的元素来得到区间和。
接下来，我们需要使用归并排序的方式对 sum 数组进行排序。在归并排序的过程中，我们需要统计满足区间和在 [lower, upper] 之间的个数。具体的做法是，对于每一个 sum 数组中的元素 sum[i] ，我们需要找到在 [lower, upper] 范围内的区间和个数。为了实现这一点，我们可以使用两个指针 windowL 和 windowR ，分别指向归并排序中左半部分和右半部分的区间。通过比较 sum[i] - upper 和 sum[i] - lower ，我们可以确定 windowL 和 windowR 的位置。然后，我们可以计算 windowR - windowL ，并将其累加到结果变量 ans 中。
最后，我们需要将归并排序后的结果拷贝回原数组 arr 。这样，我们就可以得到最终的结果。

代码实现

package class05;
 public class Code01_CountOfRangeSum {
    public static int countRangeSum(int[] nums, int lower, int upper) {
        if (nums == null || nums.length == 0) {
            return 0;
        }
        long[] sum = new long[nums.length];
        sum[0] = nums[0];
        for (int i = 1; i < nums.length; i++) {
            sum[i] = sum[i - 1] + nums[i];
        }
        return process(sum, 0, sum.length - 1, lower, upper);
    }
     private static int process(long[] sum, int l, int r, int lower, int upper) {
        if (l == r) {
            return sum[l] >= lower && sum[r] <= upper ? 1 : 0;
        }
        int m = (l + r) >>> 1;
        return process(sum, l, m, lower, upper) + process(sum, m + 1, r, lower, upper) + merge(sum, l, m, r, lower, upper);
    }
     public static int merge(long[] arr, int L, int M, int R, int lower, int upper) {
        int ans = 0;
        int windowL = L;
        int windowR = L;
        // [windowL, windowR)
        for (int i = M + 1; i <= R; i++) {
            long min = arr[i] - upper;
            long max = arr[i] - lower;
            while (windowR <= M && arr[windowR] <= max) {
                windowR++;
            }
            while (windowL <= M && arr[windowL] < min) {
                windowL++;
            }
            ans += windowR - windowL;
        }
        long[] help = new long[R - L + 1];
        int i = 0;
        int p1 = L;
        int p2 = M + 1;
        while (p1 <= M && p2 <= R) {
            help[i++] = arr[p1] <= arr[p2] ? arr[p1++] : arr[p2++];
        }
        while (p1 <= M) {
            help[i++] = arr[p1++];
        }
        while (p2 <= R) {
            help[i++] = arr[p2++];
        }
        for (i = 0; i < help.length; i++) {
            arr[L + i] = help[i];
        }
        return ans;
    }
}

测试样例

我们可以使用以下测试样例来验证我们的解决方案：

public static void main(String[] args) {
    int[] nums = { -2, 5, -1 };
    int lower = -2;
    int upper = 2;
    int result = Code01_CountOfRangeSum.countRangeSum(nums, lower, upper);
    System.out.println(result); // 输出 3
}

在这个例子中，给定数组 nums 为 [-2, 5, -1] ， lower 为 -2 ， upper 为 2 。我们需要计算区间和在 [lower, upper] 之间的个数。根据题目要求，我们可以得到区间和在 [-2, 2] 之间的个数为 3 。因此，我们的解决方案应该输出 3 。