题目链接：区间和的个数
给定一个整数数组 nums，返回区间和在 [lower, upper] 之间的个数，包含 lower 和 upper。
区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。
说明:
最直观的算法复杂度是 O(n2) ，请在此基础上优化你的算法。
示例:
输入: nums = [-2,5,-1], lower = -2, upper = 2,
输出: 3
解释: 3个区间分别是: [0,0], [2,2], [0,2]，它们表示的和分别为: -2, -1, 2。
思路
(树状数组) $O(nlogn)$
这道题一个很直观的做法就是记录前缀和，然后使用双层循环遍历所有的区间，时间复杂度$O(n^2)$。我们考虑如何来优化这个这个双层循环，我们在固定子数组的右边界的时候，采用遍历的方式求出所有区间和在[lower,upper]之间的数组个数，我们可以以更优的方式求解所有可行的区间。假设右区间为A[j]，前缀和为preSum[j]，其实我们需要求的是所有preSum[j] - upper <= preSum[i] <= preSum[j] - lower(i<j)的个数。求某一个区间内的个数，我们可以使用树状数组或者线段树来求解。
我们每次读到一个数，先把合法区间内前缀和的个数求出来（区间查询），然后将当前前缀和出现的次数加上一（单点更新）。因为只需要上面两个操作，所以可以使用树状数组来减少代码难度。整体代码思路如下：
- 求出数组的前缀和数组（包括0），并将前缀和数组离散化。
- 使用三个哈希表，分别记录每一个前缀和离散化后的大小，以该数字为右边界对应的左边界前缀和的区间
[preSum[j] - upper,preSum[j] - lower]对应的区间左右端点离散化后的值。
lower_bound查找大于等于preSum[j] - upper最小值对应的下标；
upper_bound查找大于preSum[j] - lower 最小值对应的下标；
- 先将
0对应的位置加上1（表示一个元素都没有）
- 然后遍历所有的前缀和，执行区间查询和单点更新操作。
时间复杂度：排序、离散化和树状数组的时间复杂度都是$O(nlogn)$。
代码
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
|
class Solution
{
public:
vector<int> tree;
int n, m;
inline int lowbit(int x)
{
return x & (-x);
}
void update(int idx, int delta)
{
while (idx <= m)
{
tree[idx] += delta;
idx += lowbit(idx);
}
}
int rangeSum(int l, int r)
{
//lower_bound/upper_bound找不到合法解是就是m + 1,所以直接返回0。
if (l == m + 1 || r == m + 1)
return 0;
int res = 0;
while (l != r)
{
res += tree[r] - tree[l];
r -= lowbit(r);
l -= lowbit(l);
}
return res;
}
int countRangeSum(vector<int> &nums, int lower, int upper)
{
n = nums.size();
if (n == 0)
return 0;
vector<long long> preSum(n + 1, 0);
unordered_map<int, int> hash, hash_lower, hash_upper;
for (int i = 0; i < n; i++)
preSum[i + 1] = preSum[i] + nums[i];
vector<long long> temp = preSum;
sort(temp.begin(), temp.end());
temp.erase(unique(temp.begin(), temp.end()), temp.end());
m = temp.size();
tree = vector<int>(m + 1, 0);
//这里需要注意一下下标，在temp数组中下标是`0～m-1`，但是为了和树状数组对应我们需要转化成`1~m`
//所以hash就加上了1。假设合法的区间是`[L,R]`，
//lower_bound找到的是大于等于L最小值的下标idx，加上1之后是`idx + 1`，但是在求区间和的时候，我们需要求的是左端点前面那个元素对应的前缀和，此时又要减去1。所以hash_lower就不变了。
//upper_bound找到的是大于R的第一个元素下标，也恰好是R对应的下标idx加上1之后的结果，所以治理hash_upper也不变了。
for (int i = 0; i < m; i++)
{
hash[temp[i]] = i + 1;
hash_lower[temp[i]] = lower_bound(temp.begin(), temp.end(), temp[i] - upper) - temp.begin();
hash_upper[temp[i]] = upper_bound(temp.begin(), temp.end(), temp[i] - lower) - temp.begin();
}
int res = 0;
update(hash[0], 1);
for (int i = 1; i <= n; i++)
{
res += rangeSum(hash_lower[preSum[i]], hash_upper[preSum[i]]);
update(hash[preSum[i]], 1);
}
return res;
}
};
|
最后修改于 2020-11-07
本作品采用
知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议进行许可。