【LeetCode-327】区间和的个数

问题

给定一个整数数组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）之内的 有效 区间和的个数。

示例

输入： 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 个有效区间和。

解答

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        this->lower = lower;
        this->upper = upper;
        int n = nums.size();
        vector<long> pre_sum(n + 1);
        for (int i = 0; i < n; i++)
            pre_sum[i + 1] = pre_sum[i] + nums[i];
        tmp.resize(n + 1);
        mergeSort(pre_sum, 0, n);
        return res;
    }
private:
    int lower, upper;
    int res = 0;
    vector<int> tmp;
    void mergeSort(vector<long>& nums, int left, int right) {
        if (left >= right) return;
        int mid = left + (right - left) / 2;
        mergeSort(nums, left, mid);
        mergeSort(nums, mid + 1, right);
        // ------ 与归并排序的区别
        int ii = left, jj = mid + 1, kk = mid + 1;
        while (ii <= mid) {
            while (jj <= right && nums[jj] - nums[ii] < lower) jj++;
            while (kk <= right && nums[kk] - nums[ii] <= upper) kk++;
            res += kk - jj;
            ii++;
        }
        // ------
        int i = left, j = mid + 1, cnt = left;
        while (i <= mid && j <= right) {
            if (nums[i] <= nums[j]) tmp[cnt++] = nums[i++];
            else tmp[cnt++] = nums[j++];
        }
        while (i <= mid) tmp[cnt++] = nums[i++];
        while (j <= right) tmp[cnt++] = nums[j++];
        for (int i = left; i <= right; i++) nums[i] = tmp[i];
    }
};

重点思路

1. 首先求出原数列的前缀和，两个点的前缀和之差为区间和（类似【LeetCode-304】二维区域和检索 - 矩阵不可变）；
2. 我们需要的是pre_sum中任意两点i<j，其中后面pre_sum[j] - pre_sum[i]在区间内。注意差的顺序；
3. 要保证i <= j这一条件，可以把pre_sum分成两个数组，从这两个数组分别取一个点计算区间和；
4. 当这两个数组均有序时，可以使用两个指针jj, kk依次遍历得到满足lower和upper的两个点，满足条件的数量为kk - jj；

通过上面的分析，我们对前缀和进行归并排序加上特殊处理即可满足题目要求。