力扣327 區間和的個數
阿新 • • 發佈:2020-07-22
題意
區間和的個數
給定一個整數陣列 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; } };
字首和的下標範圍。線上段樹中進行查詢