题意
区间和的个数
给定一个整数数组 nums,返回区间和在 [lower, upper] 之间的个数,包含 lower 和 upper。
区间和 S(i, j) 表示在 nums 中,位置从 i 到 j 的元素之和,包含 i 和 j (i ≤ j)。
说明:
最直观的算法复杂度是 O(n2) ,请在此基础上优化你的算法。
题解
最直接的思路是我们O(n^2)遍历区间和
int ret = 0;
for(int i=1;i<=n;i++){
//S[i] = S[i-1] + nums[i-1]
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;
但实际上我们发现第二个for循环是一个查找操作,如果S数组有序,我们可以优化查找操作到log(n),因此我们可以利用multisert来存储前缀和,实现自动排序
multiset<ll> S;
S.insert(0);
int ret = 0;
for(int i=0;i<n;i++){
presum += nums[i];
ret += distance(S.lower_bound(presum-upper),S.upper_bound(presum-lower));
S.insert(presum);
}
return ret;
我们也可以用线段树模拟这一过程,我们预处理前缀和(去重),得到数组sum。遍历原数组,对每一次的前缀和,我们查找第一个>=的位置(lower_bound),进行单点更新;根据第二个循环的公式我们可以查找满足条件的区间和。
class Solution {
#define ll long long
#define LO(sum,value) lower_bound(sum.begin(),sum.end(),value)-sum.begin();
#define UP(sum,value) upper_bound(sum.begin(),sum.end(),value)-sum.begin()-1;
public:
int d[4*100010];
int b[4*100010];
void update(int l,int r,int c,int s,int t,int p)
{
//对原数组区间[l,r]每个数加c;
//若将每个数变为c,直接用"="即可
if(l<=s&&t<=r)
{
b[p]+=c,
d[p]+=(t-s+1)*c;
return;
}
int m=(s+t)>>1;
if(b[p]&&s!=t)
{
b[2*p]+=b[p],b[2*p+1]+=b[p];
d[2*p]+=(m-s+1)*b[p],d[2*p+1]+=(t-m)*b[p];
b[p]=0;
}
if(l<=m) update(l,r,c,s,m,2*p);
if(r>m) update(l,r,c,m+1,t,2*p+1);
d[p]=d[2*p]+d[2*p+1];
}
int getsum(int l,int r,int s,int t,int p)
{
if(l<=s&&t<=r)
return d[p];
int m=(s+t)>>1;
int sum=0;
if(b[p])
{
b[2*p]+=b[p],b[2*p+1]+=b[p];
d[2*p]+=(m-s+1)*b[p],d[2*p+1]+=(t-m)*b[p];
b[p]=0;
}
if(l<=m)
sum+=getsum(l,r,s,m,2*p);
if(r>m)
sum+=getsum(l,r,m+1,t,2*p+1);
return sum;
}
int countRangeSum(vector<int>& nums, int lower, int upper) {
int N=nums.size();
vector<ll>sum(N+1);
if(N==0)
return 0;
sum[0]=0;
for(int i=1;i<=N;i++)
sum[i]=sum[i-1]+nums[i-1];
sort(sum.begin(),sum.end());
sum.erase(unique(sum.begin(),sum.end()),sum.end());
int n=sum.size();
int ans=0;
ll presum=0;
int index=LO(sum,presum);
update(index,index,1,0,n-1,1);
for(int i=0;i<N;i++)
{
presum+=nums[i];
int index=LO(sum,presum);
int l=LO(sum,presum-upper);
int r=UP(sum,presum-lower);
//cout<<l<<" "<<r<<endl;
if(l<=r)
ans+=getsum(l,r,0,n-1,1);
//cout<<ans<<endl;
update(index,index,1,0,n-1,1);
}
return ans;
}
};
前缀和的下标范围。在线段树中进行查找