[leetCode]327. 区间和的个数

题目

链接：https://leetcode-cn.com/problems/count-of-range-sum

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

说明:
最直观的算法复杂度是 O(n^2) ，请在此基础上优化你的算法。

示例:

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

归并排序

由于要求区间和因此可以想到前缀和数组，设前缀和数组为preSum。
问题可以转换为 preSum[j] - preSum[i] 属于[lower, upper]
如果给定两个升序排列的数组n1,n2, 尝试找出下标对(i, j), 满足: n2[j] - n1[i] 属于[lower, upper]
由于两个数组是升序排列的所以只需要在n2中维护两个指针 l, r, 一开始都指向n2的起始位置,首先考虑n1的第一个元素，l向有移动知道找到一个元素使n2[l] - n1[0] >= lower为止，这样l右侧的元素均大于n1[0] + lower, 再移动r指针，直到n2[r] - n1[0] > upper这样r左侧的元素均小于等于n1[0] + upper，所以区间[l, r)的所有下标j都满足n2[j] - n1[i] 属于[lower, upper] 。下面继续考察n1的第二个元素，由于n1是递增的所以l, r继续向右移动。 在上述过程中对于n1的每一个下标都应记录相应区间[l, r)的大小，最终就统计得到了下标对(i, j)的数量。

class Solution {
    private int lower;

    private int upper;

    public int countRangeSum(int[] nums, int lower, int upper) {
        this.lower = lower;
        this.upper = upper;
        int n = nums.length;
        long[] preSum = new long[n + 1];
        for (int i = 0; i < n; i++) {
            preSum[i + 1] = preSum[i] + nums[i];
        }
        return mergeSort(preSum, 0, preSum.length - 1);
    }

    public int mergeSort(long[] sums, int lo, int hi) {
        if (lo >= hi) {
            return 0;
        }
        int mid = lo + (hi - lo) / 2;
        int ans = 0;
        ans += mergeSort(sums, lo, mid);
        ans += mergeSort(sums, mid +  1, hi);
        ans += merge(sums, lo, mid, hi);
        return ans;
    }

    public int merge(long[] sums, int lo, int mid, int hi) {
        int i = lo, l = mid + 1, r = mid + 1;
        int count = 0;
        while (i <= mid) {
            while (l <= hi && sums[i] + lower > sums[l]) { // l 右侧的元素都大于等于 sums[i] + lower
                l++;
            }
            while (r <= hi && sums[i] + upper >= sums[r]) { // r 左侧的元素都小于等于 sums[i] + upper
                r++;
            }
            count+= r - l;
            i++;
        }

        int[] aux = new int[hi - lo + 1]; 
        int idx = 0;
        int p1 = lo, p2 = mid + 1;
        while (p1 <= mid && p2 <= hi) {
            if (sums[p1] < sums[p2]) {
                aux[idx++] = (int) sums[p1++];
            } else {
                aux[idx++] = (int) sums[p2++];
            }
        }

        while (p1 <= mid) {
            aux[idx++] = (int) sums[p1++];
        }
        while (p2 <= hi) {
            aux[idx++] = (int) sums[p2++];
        }

        for (int j = lo; j <= hi; j++) {
            sums[j] = aux[j-lo];
        }

        return count;
    }
}