暴力法
简单且超时58。。。
class Solution {
public:
int countRangeSum(vector<int>& nums, int lower, int upper) {
int count = 0;
for(int i = 0; i < nums.size(); ++i){
if(nums[i] <= upper && nums[i] >= lower) ++count;
long sum = nums[i];
for(int j = i+1; j < nums.size(); ++j){
sum += nums[j];
if(sum <= upper && sum >= lower) ++count;
}
}
return count;
}
};
那就想想怎么优化吧。一说到优化，可以存储前缀和：
class Solution {
public:
int countRangeSum(vector<int>& nums, int lower, int upper) {
int n= nums.size();
long presum = 0;
vector<long> S(n+1,0);
int ret = 0;
for(int i=1;i<=n;i++){
presum += nums[i-1];
for(int j=1;j<=i;j++){
if(lower <= presum - S[j-1] && presum - S[j-1] <= upper) ret++;
}
S[i] = presum;
}
return ret;
}
};
其实复杂度一样，只是采取了前缀和的思路来计算区间和，也是超时的。
归并
先考虑前缀和，举个例子：
nums = [4,2,3,7,-1,1,5]，lower = 0, upper = 15
前缀和数组为:
temp = [4,6,9,16,15,16,21]
我们可以利用前缀和之间的差值来表示区间和：
temp[4] - temp[1] = 15 - 6 = 9
9 = nums[2] + nums[3] + nums[4]
也就是说，有了前缀和数组，我们想要计算某个区间的和是否符合要求，就不用进行累加了，直接利用区间边界对应的前缀和进行相减就行。
参考：327.区间和的个数题解综合 - 区间和的个数 - 力扣（LeetCode）
class Solution {
public:
vector<long> tmp;//归并排序辅助数组
void merge(vector<long> &S,int left, int mid, int right){//归并排序标准merge函数了
int i = left, j = mid + 1;
int k = 0;
while(i <= mid && j <= right){
if(S[i] <= S[j]) tmp[k++] = S[i++];
else tmp[k++] = S[j++];
}
while(i <= mid) tmp[k++] = S[i++];
while(j <= right) tmp[k++] = S[j++];
copy(tmp.begin(), tmp.begin() + k, S.begin() + left);//使用标准库函数好像更快一点
}
int merge_sort(vector<long> &S,int left,int right,int lower,int upper){
if(left >= right) return 0;
int cnt = 0;
int mid = left + (right-left)/2;
cnt += merge_sort(S, left, mid, lower, upper);
cnt += merge_sort(S, mid+1, right, lower, upper);
//计算区间和符合要求的个数
int i = left;
int j = mid + 1, k = mid + 1;
while(i <= mid){//对左半部分的每个元素
//auto it1 = lower_bound(S.begin()+mid+1, S.begin() + right+1, S[i] + lower);
//auto it2 = upper_bound(S.begin()+mid+1, S.begin() + right+1, S[i] + upper);
while(k <= right && S[k] - S[i] < lower)k++;//区间和 < lower
while(j <= right && S[j] - S[i] <= upper)j++;//区间和 <= 上界
cnt += j - k;
//cnt += it2 - it1;
++i;
}
if(S[mid] <= S[mid+1]) return cnt; //加一个判断，此时数组已经是有序的，不用进行后续操作了
merge(S, left, mid, right);
return cnt;
}
int countRangeSum(vector<int>& nums, int lower, int upper) {
int n = nums.size();
vector<long> S(n+1,0);//前缀和数组
tmp = vector<long>(n+1);
for(int i=1;i<=n;i++) S[i] = S[i-1] + nums[i-1];
return merge_sort(S, 0, n, lower, upper);
}
};
平衡树
这个思路很好，借助平衡树进行排序，刚好平衡树还能二分查找，又可以使用distance函数计算距离。
class Solution {
public:
int countRangeSum(vector<int>& nums, int lower, int upper) {
int n= nums.size();
long presum = 0;
//本质是平衡二叉查找树（红黑树）
multiset<long> S;//因为有重复元素，而且前缀和并不有序，使用multiset存储并排序
S.insert(0);
int ret = 0;
//区间和转化为前缀和的相减
for(int i=0;i<n;i++){
presum += nums[i];//一边插入一边查找
//不过对于红黑树而言，distance函数是O(n)的
//当然distance函数对于红黑树来说，可以简单地一直递增迭代器得到
ret += distance(S.lower_bound(presum-upper), S.upper_bound(presum-lower));//二分查找上下界
S.insert(presum);
}
return ret;
}
};