327. 区间和的个数

题目描述

给定一个整数数组 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。

思路

暴力

计算全部区间，$O(n^2)$，超时。

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size(), res = 0;
        for(int i = 0;i < n;++i) {
            long long rangesum = 0L;
            for(int j = i;j < n;++j) {
                rangesum += nums[j];
                if(rangesum >= lower && rangesum <= upper) ++res;
            }
        }
        return res;
    }
};

分治

前缀和

计算前缀和presum，问题转化为对任意的前缀和val = presum[i] - presum[j](i >= j)落在 [lower, upper] 之间的个数。

计算前缀和注意第一个元素要补0。

若前缀和无序，计算的时间复杂度依旧是$O(n^2)$，若前缀和有序。计算时，可以使用双指针。

问题划分

若两个区间[l,m)，[m, r)的前缀和都是有序的，可以利用三指针。

• p0指向l， p1指向m，p2指向m

• p1后移直到presum[p1] - presum[p0] >= lower

• p2后移直到presum[p2] - presum[p0] > upper

• p2 - p1则为以p0为起点的区间数量，由于[l, m)是有序的，因此p0后移之后，p1，p2必定后移，不用向前回溯，减少了时间。

代码

class Solution {
public:
    int merage(vector<long long>& presum, int l, int r, int lower, int upper) {
        // 左右右开，[l,r)，区间长度为1时，直接返回0
        if(l == r - 1 || l == r) return 0;
        int m = l + (r - l)/2;
        int res = 0;
        res += merage(presum, l, m, lower, upper);
        res += merage(presum, m, r, lower, upper);
        int p0 = l, p1 = m, p2 = m;
        while(p0 < m) {
            while(p1 < r && presum[p1] - presum[p0] < lower) ++p1;
            while(p2 < r && presum[p2] - presum[p0] <= upper) ++p2;
            res += p2 - p1;
            ++p0;
        }
        vector<long long> temp(r - l);
        for(int i = 0, j = l, k = m;i < r - l;++i) {
            if(j < m && k < r) {
                if(presum[j] < presum[k]) 
                temp[i] = presum[j++];
                else 
                temp[i] = presum[k++];
            } else if(j < m){
                temp[i] = presum[j++];
            } else {
                temp[i] = presum[k++];
            }
        }
        for(int i = 0;i < r - l;++i) {
            presum[i + l] = temp[i];
        }
        return res;
    }
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size();
        vector<long long> presum(n + 1);
        for(int i = 0;i < nums.size();++i) {
            presum[i+1] = presum[i] + nums[i];
        }
        return merage(presum, 0, n + 1, lower, upper);
    }
};

时间复杂度

$T(N) = 2*T(N/2) + O(N)$

$T(N) = O(NlogN)$

空间复杂度

设空间占用为 $M(N)$，递归调用所需空间为 $M(N/2)$，而合并数组所需空间为 O(N)O(N)，故

$$M(N) = \max\big\{M(N/2), O(N)\big\} = M(N/2) + O(N)$$

$M(N) = O(N)$