[LeetCode] Count of Range Sum | 区间和统计

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

带条件约束的统计。结果 = 二分后左边的统计 + 右边的统计 + 横跨中间的，难点在于横跨中间的计算。

这个问题和“数组中的逆序对”类似，只不过对统计的条件要求更一般化，前者可以用线段树、树状数组、归并排序（排序的目的是为了加速“横跨中间的统计”的计算速度），本问题也一样。

设S(i,j) = nums[i] + ... + nums[j] = S(j) - S(i-1)表示原数组的一个区间和(i <= j)，现要找出所有满足lower <= S(i,j) <= upper的区间，即S(i-1)+lower <= S(j) <= S(i-1) + upper。

计算出所有的S(i)后，在归并排序的框架中，先分治计算左、右两边的统计值，然后在合并的过程中计算“横跨中间的”（也可以之后合并）。如果按从小到大排序，就遍历左边的S[i]，对每一个S[i]，在右边寻找在范围[S(i)+lower, S(i) +upper]之间的所有元素的数量，由于是有序数组的查找，因而可以快速完成。

另外就是要注意计算S[i]的时候，最好用长整型以避免可能存在的整数溢出。

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        if (nums.empty()) return 0;
        vector<long long> S(nums.size());
        init(nums, S, lower, upper);
        return count(S, 0, nums.size() - 1);
    }
    
private:
    int _lower, _upper;
    
    void init(vector<int>& nums, vector<long long>& S, int lower, int upper) {
        _lower = lower;
        _upper = upper;
        
        S[0] = nums[0];
        for (int i = 1; i < (int) nums.size(); ++i) {
            S[i] = S[i-1] + nums[i];
        }
    }
    
    int count(vector<long long>& S, int start, int end) {
        if (start > end) return 0;
        if (start == end) return (S[start] >= _lower && S[start] <= _upper? 1:0);
        
        int mid = start + (end - start) / 2;
        int left = count(S, start, mid);
        int right = count(S, mid + 1, end);
        int cross = 0;
        vector<long long>::iterator low, up;
        for (int i = start; i <= mid; ++i) {
            low = lower_bound(S.begin() + mid+1, S.begin() + end+1, S[i] + _lower);
            up = upper_bound(S.begin() + mid+1, S.begin() + end+1, S[i] + _upper);
            cross += (up - low);
        }
        merge(S, start, mid, end);
        
        return ( left + right + cross );
    }
    
    void merge(vector<long long>& S, int start, int mid, int end) {
        int n = end - start + 1;
        long long helper[n];
        int p = start, q = mid + 1, r = 0;
        while (p <= mid && q <= end) {
            if (S[p] < S[q]) helper[r++] = S[p++];
            else helper[r++] = S[q++];
        }
        while (p <= mid) helper[r++] = S[p++];
        while (q <= end) helper[r++] = S[q++];
        copy(helper, helper+n, S.begin()+start);
    }
};

相似问题:

• https://leetcode.com/problems/count-of-smaller-numbers-after-self/#/description
• https://leetcode.com/problems/reverse-pairs/#/description
• 重要的逆序对,同上: https://acm.sjtu.edu.cn/OnlineJudge/problem/1219