1. 程式人生 > 實用技巧 >力扣327 區間和的個數

力扣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;     
    }
};

字首和的下標範圍。線上段樹中進行查詢