[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);
}
};
相似问题: