题目描述
给定一个整数数组 nums
,返回区间和在 [lower, upper]
之间的个数,包含 lower
和 upper
。
区间和 S(i, j)
表示在 nums
中,位置从 i
到 j
的元素之和,包含 i
和 j
(i
≤ j
)。
说明
- 最直观的算法复杂度是 $O(n^2)$ ,你必须使用更好的算法。
样例
输入:nums = [-2,5,-1], lower = -2, upper = 2,
输出:3
解释:3个区间分别是: [0,0], [2,2], [0,2],它们表示的和分别为: -2, -1, 2。
算法1
(树状数组) $O(n \log n)$
- 我们考虑前缀和数组。对于每个位置
i
的前缀和,我们累计i
之前范围在[sum - upper, sum - lower]
的前缀和的个数。 - 用高级数据结构的实现方式有两种,一是采用平衡二叉搜索树,每次将一个前缀和放入二叉搜索树,然后直接查询。二是将所有有关的数字离散化后,用树状数组求出区间和。
- 这里我们采用第二种实现,首先将 0,每个位置的前缀和,以及每个位置前缀和
sum - upper
和sum - lower
放入到离散化数组中。然后排序去重离散化数组,得到每个值的离散化后的值,范围在[1, m]
。 - 我们开一个范围从
[1, m]
的树状数组,更新 0 所在离散化值的位置加 1。然后按照第一步的逻辑统计答案,最后将这个位置的离散化值的位置在树状数组中加 1。
时间复杂度
- 离散化的时间复杂度为 $O(n \log n)$,树状数组单次操作的时间复杂度为 $O(\log n)$。
- 故总时间复杂度为 $O(n \log n)$。
空间复杂度
- 需要额外的离散化数组和树状数组,故空间复杂度为 $O(n)$。
C++ 代码
#define LL long long
class Solution {
public:
int m;
vector<int> f;
void update(int x, int y) {
for (; x <= m; x += x & -x)
f[x] += y;
}
int query(int x) {
int tot = 0;
for (; x; x -= x & -x)
tot += f[x];
return tot;
}
int countRangeSum(vector<int>& nums, int lower, int upper) {
int n = nums.size();
vector<LL> pre;
pre.push_back(0);
LL sum = 0;
for (int i = 0; i < n; i++) {
sum += nums[i];
pre.push_back(sum);
pre.push_back(sum - upper);
pre.push_back(sum - lower);
}
sort(pre.begin(), pre.end());
m = unique(pre.begin(), pre.end()) - pre.begin();
pre.resize(m);
f.resize(m + 1, 0);
int zero = lower_bound(pre.begin(), pre.end(), 0) - pre.begin() + 1;
update(zero, 1);
int ans = 0;
sum = 0;
for (int i = 0; i < n; i++) {
sum += nums[i];
LL x = lower_bound(pre.begin(), pre.end(), sum - upper) - pre.begin() + 1;
LL y = lower_bound(pre.begin(), pre.end(), sum - lower) - pre.begin() + 1;
ans += query(y) - query(x - 1);
LL num = lower_bound(pre.begin(), pre.end(), sum) - pre.begin() + 1;
update(num, 1);
}
return ans;
}
};
算法2
(分治) $O(n \log n)$
- 我们采取分治的思想。首先求出前缀和数组(包括开头的 0)
sum
,在这个数组上采用分治算法:每次将数组平均分为两部分,递归处理左右两部分内部的答案,然后将左右两部分内部从小到大排序,最后归并跨越左右两部分的答案。 - 接下来讨论如何求跨越两部分的答案:假设左右两部分已经从小到大排好序,我们设 $i$ 为右部分的某个位置,对于每个 $i$,都有一个在左部分连续的区间
[j, k]
, 对应着合法答案的区间,即sum[i] - sum[j]
,sum[i] - sum[j + 1], ..., sum[i] - sum[k]
都是在[lower, upper]
中。随着 $i$ 向右移动,这个区间也会整体向右移动。这给我们提供了滑动窗口的思想,我们实时维护这个窗口,可以在均摊 $O(1)$ 的时间内找到每个 $i$ 对应的区间。 - 然后利用归并排序的思想将当前区间排序。
时间复杂度
- 时间和归并排序一样,每一层内部的时间复杂度为 $O(n)$,共 $O(\log n)$ 层,故总时间复杂度为 $O(n \log n)$。
空间复杂度
- 归并排序需要线性的额外空间。
C++ 代码
#define LL long long
class Solution {
public:
int ans;
void solve(vector<LL>& sum, int l, int r, int lower, int upper) {
if (l == r)
return;
int mid = (l + r) >> 1;
solve(sum, l, mid, lower, upper);
solve(sum, mid + 1, r, lower, upper);
for (int i = mid + 1, j = l, k = l; i <= r; i++) {
while (j <= mid && sum[i] - sum[j] > upper) j++;
while (k <= mid && sum[i] - sum[k] >= lower) k++;
ans += k - j;
}
vector<LL> cpy(r - l + 1);
int cnt = 0;
int i = l, j = mid + 1;
while (i <= mid || j <= r) {
if (i == mid + 1) cpy[cnt++] = sum[j++];
else if (j == r + 1) cpy[cnt++] = sum[i++];
else {
if (sum[i] < sum[j]) cpy[cnt++] = sum[i++];
else cpy[cnt++] = sum[j++];
}
}
for (int i = l; i <= r; i++) sum[i] = cpy[i - l];
}
int countRangeSum(vector<int>& nums, int lower, int upper) {
int n = nums.size();
vector<LL> sum(n + 1);
sum[0] = 0;
for (int i = 1; i <= n; i++)
sum[i] = sum[i - 1] + nums[i - 1];
ans = 0;
solve(sum, 0, n, lower, upper);
return ans;
}
};
大佬昨天研究的时候 有一点没有想清楚,就是为啥要把0的位置 初始化为1 不初始化会怎么样呢
0 代表开始的位置,即前缀和本身
明白了 hh
解法一离散化后范围应该在[0, m -1]吧,然后f数组多开一位相当于把所有前缀后移一位方便求区间和,不知道我这样理解对不对
可以这么理解
谢大佬~
太强了