给你一个整数数组 nums 以及两个整数 lower 和 upper 。求数组中,值位于范围 [lower, upper] (包含 lower 和 upper)之内的 区间和的个数 。
区间和 S(i, j) 表示在 nums 中,位置从 i 到 j 的元素之和,包含 i 和 j (i ≤ j)。
示例 1:
输入:nums = [-2,5,-1], lower = -2, upper = 2
输出:3
解释:存在三个区间:[0,0]、[2,2] 和 [0,2] ,对应的区间和分别是:-2 、-1 、2 。
示例 2:
输入:nums = [0], lower = 0, upper = 0
输出:1
提示:
1 <= nums.length <= 105
-231 <= nums[i] <= 231 - 1
-105 <= lower <= upper <= 105
题目数据保证答案是一个 32 位 的整数
解题思路:
前缀和 + 归并排序。对前缀和进行 排序,在排序的途中求出满足要求的和的数量。
代码如下:
class Solution {
public:
int countRangeSum(vector<int>& nums, int lower, int upper) {
int n = nums.size();
vector<long>sum(n, 0);
sum[0] = nums[0];
for (int i = 1; i < n; ++i) sum[i] = sum[i - 1] + nums[i];
return helper(sum, lower, upper, 0, n);
}
int helper(vector<long>&sum, int lower, int upper, int left, int right) {
if (left >= right) return 0;
else if (right - left == 1) return (sum[left] >= lower && sum[left] <= upper);
int mid = left + (right - left) / 2;
int left_count = helper(sum, lower, upper, left, mid);
int right_count = helper(sum, lower, upper, mid, right);
int i = left;
int count = left_count + right_count;
vector<long>cache(right - left,0);
int l = mid, r = mid, k = 0, j = mid;
while (i < mid) {
while (l < right && sum[l] - sum[i] < lower )l++;
while (r < right && sum[r] - sum[i] <= upper)r++;
count += (r - l);
while (j < right && sum[j] < sum[i]) cache[k++] = sum[j++];
cache[k++] = sum[i++];
}
while (j < right) cache[k++] = sum[j++];
for (int pos = 0; left + pos < right; ++pos) sum[left + pos] = cache[pos];
return count;
}
};