HackerEarth Good Evening sweetheart!![Medium] solution

In this HackerEarth Good Evening sweetheart!![Medium] problem solution Scube sold you N candles, for each candle i you know its height. Using a candle during one evening decreases the candle height by 1. You plan to have at most M romantic evenings. For each evening i you know the number of candles C[i] you want to lit. Find a strategy of lighting the candles in order to maximize the number of evenings you can spend. You are forced to stop after the night i when you can’t light C[i] candles.

HackerEarth Good Evening sweetheart!![Medium] problem solution.

include<bits/stdc++.h>
using namespace std;
#define int long long
#define rep(i,n) for(int i=0;i<(n);i++)
#define pb push_back
#define all(v) (v).begin(),(v).end()
#define fi first
#define se second
typedef vector<int>vint;
typedef pair<int,int>pint;
typedef vector<pint>vpint;
template<typename A,typename B>inline void chmin(A &a,B b){if(a>b)a=b;}
template<typename A,typename B>inline void chmax(A &a,B b){if(a<b)a=b;}
struct segtree{
    static const int SEG=1<<17;
    vint dat,put;
    segtree():dat(SEG*2),put(SEG*2){}
    inline void push(int k,int l,int r){
        dat[k]+=put[k]*(r-l);
        if(k<SEG-1){
            put[k*2+1]+=put[k];
            put[k*2+2]+=put[k];
        }
        put[k]=0;
    }
    void update(int a,int b,int x,int k=0,int l=0,int r=SEG){
        push(k,l,r);
        if(r<=a|b<=l)return;
        if(a<=l&&r<=b){
            put[k]+=x;
            push(k,l,r);
            return;
        }
        update(a,b,x,k*2+1,l,(l+r)/2);
        update(a,b,x,k*2+2,(l+r)/2,r);
        dat[k]=dat[k*2+1]+dat[k*2+2];
    }
    int query(int a,int b,int k=0,int l=0,int r=SEG){
        push(k,l,r);
        if(r<=a||b<=l)return 0;
        if(a<=l&&r<=b)return dat[k];
        return query(a,b,k*2+1,l,(l+r)/2)
        +query(a,b,k*2+2,(l+r)/2,r);
    }
};
int N,M;
int H[111111],C[111111];
signed main(){
    cin>>N>>M;
    rep(i,N)cin>>H[i];
    rep(i,M)cin>>C[i];
   sort(H,H+N);
    segtree seg;
    rep(i,N)seg.update(i,i+1,H[i]);
    rep(i,M){
        int x=seg.query(N-C[i],N-C[i]+1);
        int lb=-1,ub=N-C[i];
        while(ub-lb>1){
            int mid=(ub+lb)/2;
            if(seg.query(mid,mid+1)==x)ub=mid;
            else lb=mid;
        }
        int l=ub;

        lb=N-C[i],ub=N;
        while(ub-lb>1){
            int mid=(ub+lb)/2;
            if(seg.query(mid,mid+1)==x)lb=mid;
            else ub=mid;
        }
        int r=ub;
        if(x==0){
            cout<<i<<endl;
            return 0;
        }
        seg.update(r,N,-1);
        seg.update(l,l+C[i]-(N-r),-1);
    }
    cout<<M<<endl;
}

coding problems