327. Count of Range Sum

问题：

给定一个数组，求连续元素之和在给定范围[lower, upper]之间的，连续idx为(i～j)元素组个数。

Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.

Example:
Input: nums = [-2,5,-1], lower = -2, upper = 2,
Output: 3 
Explanation: The three ranges are : [0,0], [2,2], [0,2] and their respective sums are: -2, -1, 2.
 

Constraints:
0 <= nums.length <= 10^4

　　

解法：

解法类似  315. Count of Smaller Numbers After Self

解法一：FenwickTree

• 只是需要多求一步presum->sum，作为处理对象。
• 对前缀和数组，进行sort->sortsum，来确定在FenwickTree中的idx。
• FenwickTree记录对应相同sum值的个数（出现频率）。

然后对原顺序的sum，逐一读入，

每个sum[i]中，在FenwickTree中，寻找sum[i]-upper ~ sum[i]-lower 范围内共有多少个，res追加。

然后把当前sum[i]找到idx，插入FenwickTree中，累计count+1。

 

代码参考：

class FenwickTree {
public:
    FenwickTree(int n):tree(n+1, 0) {}
    void update(int i, int delta) {
        while(i<tree.size()) {
            tree[i]+=delta;
            i+=lowbit(i);
        }
    }
    int getPresum(int i) {
        int sum=0;
        while(i>0){
            sum+=tree[i];
            i-=lowbit(i);
        }
        return sum;
    }
private:
    vector<int> tree;
    int lowbit(int x) {
        return x&(-x);
    }
};

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int cursum=0;
        int res=0;
        vector<long long> sum(nums.size()+1,0);
        for(int i=0; i<nums.size(); i++){
            sum[i+1]=sum[i]+nums[i];
        }
        vector<long long> sortsum(sum);
        sort(sortsum.begin(), sortsum.end());
        
        FenwickTree ftree(sum.size());
        for(int i=0; i<sum.size(); i++){
            int idx_low = distance(sortsum.begin(), lower_bound(sortsum.begin(),sortsum.end(),sum[i]-upper));
            int idx_upper = distance(sortsum.begin(), upper_bound(sortsum.begin(),sortsum.end(),sum[i]-lower));
            res+=(ftree.getPresum(idx_upper)-ftree.getPresum(idx_low));
            
            int idx = distance(sortsum.begin(), lower_bound(sortsum.begin(),sortsum.end(),sum[i]));
            ftree.update(idx+1, 1);
        }
        return res;
    }
};

 

解法二：mergeSort

同上，先求得presum->sum，作为处理对象。

在mergeSort的同时，

对已经排序好的前后两个数列

• start ~ mid
• mid+1 ~ end

遍历第一个数列（start ~ mid）的每一个sum[p]

在后一个数列（mid+1 ~ end）中，

• 找sum[m]，（使sum[m]不断增大）一直找到使得：sum[m]-sum[p] >= lower
• 找sum[n]，（使sum[n]不断增大）一直找到使得：sum[n]-sum[p] > upper

若不满足（sum[m]-sum[p] < lo   sum[n]-sum[p] <= up），则一直++，直到找到停止while循环。

那么m-n的个数=要求的对于sum[p]满足条件的个数。

        int m=mid+1, n=mid+1;
        for(int p = start; p <= mid; p++) {
            while(m<=end && sum[m]-sum[p]<lo) m++;
            while(n<=end && sum[n]-sum[p]<=up) n++;
            res += (n-m);
        }

♻️ 优化：这里随着下一个sum[p]是增大的，那么上次找到的临界值m和n，不需要从头再算一次，直接从上一次的临界值继续尝试增大即可。

累计后，将两个数列merge排序。

 

代码参考：

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int res=0;
        lo = lower;
        up = upper;
        sum.resize(nums.size()+1,0);
        for(int i=0; i<nums.size(); i++){
            sum[i+1]=sum[i]+nums[i];
        }
        res = mergeSort(0, sum.size()-1);
        return res;
    }
    int mergeSort(int start, int end) {
        int res = 0;
        if(end<=start) return res;
        int mid = start + (end - start) / 2;
        res += mergeSort(start, mid);
        res += mergeSort(mid+1, end);
        res += merge(start, mid, end);
        return res;
    }
    int merge(int start, int mid, int end) {
        int res = 0;
        int m=mid+1, n=mid+1;
        for(int p = start; p <= mid; p++) {
            while(m<=end && sum[m]-sum[p]<lo) m++;
            while(n<=end && sum[n]-sum[p]<=up) n++;
            res += (n-m);
        }
        vector<long long> tmp(end-start+1, 0);
        int p=start, q=mid+1;
        int i=0;
        while(p<=mid && q<=end){
            if(sum[p]<sum[q]) tmp[i++] = sum[p++];
            else tmp[i++] = sum[q++];
        }
        while(p<=mid) tmp[i++] = sum[p++];
        while(q<=end) tmp[i++] = sum[q++];
        copy(tmp.begin(), tmp.end(), sum.begin()+start);
        
        return res;
    }
private:
    int lo;
    int up;
    vector<long long> sum;
};