题目
暴力解法
就如题目所说最直观的方法是O(n * n)。
- 使用数组保存前缀和，时间复杂度O(n)
- 两遍for循环，找到每段区间的和，并判断是否在范围内，时间复杂度O(n * n)
代码
class Solution {
public:
//暴力解法
int countRangeSum(vector<int>& nums, int lower, int upper) {
long long n = nums.size();
long long sum = 0;
vector<long long> dp;
//求前缀和
for (long long i = 0; i < n; i++) {
sum += nums[i];
dp.push_back(sum);
}
long long res = 0, cur;
for (long long i = 0; i < n; i++) {
for (long long j = i; j < n; j++) {
cur = i == 0 ? dp[j] : dp[j] - dp[i - 1];
//判断范围
if (cur >= lower && cur <= upper) {
res++;
}
}
}
return res;
}
};
树状数组
概念：树状数组(Binary Indexed Tree(B.I.T), Fenwick Tree)是一个查询和修改复杂度都为log(n)的数据结构。主要用于查询任意两位之间的所有元素之和，但是每次只能修改一个元素的值；经过简单修改可以在log(n)的复杂度下进行范围修改，但是这时只能查询其中一个元素的值(如果加入多个辅助数组则可以实现区间修改与区间查询)。
学习的话，看这篇B站资源。
通过上述，我们已经知道查询一个区间的前缀和，只需要O(log(n))的时间复杂度，这道题主要利用的就是这个性质。
步骤(通过一个例子来说明: nums = [-2,5,-1], lower = -2, upper = 2)
第一步仍然是求出前缀和数组，sum = [0, -2, 3, 2]。之所以第一位多加一个0，是为了方便计算第0位到i位的前缀和。
有了前缀和数组，我们取前缀和数组里两个不同位置的值a和b，其中a是0到i位置的和，b是0到j位置的和，i < j 。
所求的结果可以转化成：有多少组 lower <= b - a <= upper。
上述表达式可以化为 lower + a <= b <= upper + a，即对于每一个a，有多少个b存在于[lower + a, upper + a]区间内。
但是对于上述关系，我们并不需要知道lower + a, b, upper + a三个具体的值是多少，我们只需要他们的相对关系。所以我们可以把所有的数据（包括a，a + lower，a + upper）离散化。
sum数组添加上每个元素和lower，upper的和之后变为。[0, -2, 3, 2, -2, 2, -4, 0, 1, 5, 0, 4]，再离散化的结果为[3, 2, 6, 5, 2, 5, 1, 3, 4, 8, 3, 7]
然后对原sum数组逆序遍历，首先遍历到a = 2，lower + a = 0，0离散化后为3， upper + a = 4， 4离散化后为7。
通过树状数组查询4到7之间的和，通过query(7) - query(4 - 1)得到。将查询结果加到最终结果上。再将树状数组里，a = 2对应的离散化值5的计数加一。
对原sum数组每一个数做前两步的操作，就能得到最后的统计结果。
时间复杂度：求前缀和数组 + 离散化前缀和数组 + 遍历前缀和数组 * （2 * 树状数组查询 + 树状数组单点修改) = O(n) + O(n * log(n)) + O(n * (2 * log(n) + log(n))) = O(n * log(n))
代码
class NumTree{
private:
vector<int> vec;
int n;
int lowbit(int x) {
return x & -x;
}
public:
NumTree(int size){
n = size;
vec = vector<int>(n + 1, 0);
}
void change(int x) {
for (; x <= n; x += lowbit(x)) {
vec[x]++;
}
}
int query(int x) {
int res = 0;
for (; x; x -= lowbit(x)) {
res += vec[x];
}
return res;
}
int query_area(int i, int j) {
return query(j) - query(i - 1);
}
};
class Solution {
public:
int countRangeSum(vector<int>& nums, int lower, int upper) {
if (nums.empty()){
return 0;
}
int n = nums.size();
//前缀合
long long cur = 0;
vector<long long> sum;
sum.push_back(0);
for (auto & x : nums) {
cur += x;
sum.push_back(cur);
}
//去重
set<long long> st;
for (int i = 0; i <= n; i++) {
st.insert(sum[i]);
st.insert(sum[i] + lower);
st.insert(sum[i] + upper);
}
//离散化
unordered_map<long long, int> mp;
int cnt = 0;
for (auto & x : st) {
mp[x] = ++cnt;
}
NumTree numTree(cnt + 1);
int res = 0;
for (int i = n; i >= 0; i--) {
res += numTree.query_area(mp[sum[i] + lower], mp[sum[i] + upper]);
numTree.change(mp[sum[i]]);
}
return res;
}
};
