题目描述

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

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

示例:

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

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

解法1

先解释下naive的做法,为了计算区间和,一般来说prefix sum是少不了的。有了prefix sum,我们可以在O(1)的时间计算出区间和。首先定义一个数组preSum[i] = sum(nums[0]…nums[i])。然后我们计算符合lower<=preSum[j] – preSum[i]<=upper的次数。显然,这种做法的时间复杂度是O(n^2)。

解法1采用merge sort。我们回忆“315.计算右侧小于当前元素的个数”,我们将题目的目标形式化定义:

counts[i] = 统计 nums[i] < nums[j] 成立的次数,0 <= i < j < n

在那道题中,我们将nums进行merge sort,在merge的过程时,用rightCount变量统计nums[i] > nums[j]的次数,当下一次nums[i] <= nums[j]成立时,将rightCount写入counts[i]。

如果我们将这道题的目标形式化定义,可以得到与315题相似的形式。

count[i] = 统计 lower <= preSum[j] - preSum[i] <= upper 成立的次数, i < j
answer = sum(count[0]...count[n-1])

这道题,我们同样可以采用merge sort方法。在merge的过程中将preSum分为左右两部分,他们的下标范围分别是[start, mid)与[mid, end)。我们在左区间寻找索引i ∈ [start, mid),在右区间[mid, end)寻找第一个使得preSum[j] – preSum[i] >= lower, preSum[k] – preSum[i] > upper成立的j, k ,其中j<k且j, k ∈ [mid, end)。

如果我们找到了这样一对i, j, k,我们就找到了k – j个区间和落在[lower, upper]的范围。想想这是为什么呢?首先,preSum[i]存储的是sum(nums[0]…nums[i]),那么preSum[j] – preSum[i] = sum(nums[i+1]…nums[j]). 图1说明了利用preSum计算nums[3…4]的区间和。

图1 使用preSum计算区间和

如果preSum[j] – preSum[i] >= lower, 那么preSum[j + 1] – preSum[i] >= lower也成立,这是因为preSum[start…mid)与preSum[mid…end)已经被排过序了,使得preSum[j+1] >= preSum[j]成立。如果preSum[j] – preSum[i] >= lower, preSum[k] – preSum[i] > upper,那么我门就找到一些p,j<=p<k,有lower<=sum(nums[i+1]…nums[p])<=upper成立,我们就找到k-j个p使得sum(nums[i+1]…nums[p])落到给定的区间。

我在merge sort的merge操作之前,在[start, mid)寻找i,在[mid, end)寻找j,k,对于每一个i,找到j-k个区间和落在[lower, upper],我们将这些区间和累积即可。需要注意一点,当end – start <=1,说明我们只有一个元素,我们直接判断这个元素是否落到[lower, upper]即可。

你可能会思考,为什么在merge sort的左半部分寻找i,在右半部分间寻找j, k,而不是在左右分别找出i, j, k?如果你在左右两侧独立的寻找i, j, k,那么你就会落掉跨越[start, mid), [mid, end)的区间和。

我们在编写代码时,一定要想好变量的定义,例如是取[start, end]还是[start, end),不要在一个函数用闭区间写法,另一个函数用左闭右开写法,否则会迷失在思考边界的大海里,另bug层出不穷……。全部代码如下所示,时间复杂度O(nlogn),空间复杂度O(n)。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        if (nums == null || nums.length == 0)
            return 0;
        int n = nums.length;
        long[] preSum = new long[n];

        preSum[0] = nums[0];
        for (int i = 1; i < n; i++) {
            preSum[i] = preSum[i - 1] + nums[i];
        }

        // [start, end)
        return mergeSort(preSum, 0, n, lower, upper);
    }

    int mergeSort(long[] preSum, int start, int end, int lower, int upper) {
        int mid = (start + end) >>> 1;

        if (end - start <= 1)
            return preSum[start] >= lower && preSum[start] <= upper ? 1 : 0;

        int count = mergeSort(preSum, start, mid, lower, upper) +
                mergeSort(preSum, mid, end, lower, upper);

        int j = mid, k = mid;

        for (int i = start; i < mid; i++) {
            while (j < end && preSum[j] - preSum[i] < lower)
                j++;

            while (k < end && preSum[k] - preSum[i] <= upper)
                k++;

            count += k - j;
        }

        merge(preSum, start, end);

        return count;
    }

    void merge(long[] preSum, int start, int end) {
        int mid = (start + end) >>> 1;
        int i = start, j = mid;
        long[] sorted = new long[end - start];
        int sortedIdx = 0;

        while (i < mid && j < end) {
            if (preSum[i] <= preSum[j]) {
                sorted[sortedIdx++] = preSum[i++];
            } else {
                sorted[sortedIdx++] = preSum[j++];
            }
        }

        while (i < mid) {
            sorted[sortedIdx++] = preSum[i++];
        }

        while (j < end) {
            sorted[sortedIdx++] = preSum[j++];
        }

        for (sortedIdx = 0; sortedIdx < sorted.length; sortedIdx++)
            preSum[start + sortedIdx] = sorted[sortedIdx];
    }
}

参考:https://leetcode.com/problems/count-of-range-sum/discuss/77990/Share-my-solution

pwrliang Algorithms, Sort ,

Leave a Reply

Your email address will not be published. Required fields are marked *