力扣327 区间和的个数

题意

区间和的个数
给定一个整数数组 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;     
    }
};

前缀和的下标范围。在线段树中进行查找

posted @ 2020-07-22 12:24  blueattack  阅读(229)  评论(0编辑  收藏  举报