Submission #1138308

#	Time	Username	Problem	Language	Result	Execution time	Memory
1138308		imarn	Paths (RMI21_paths)	C++20	100 / 100	452 ms	39352 KiB

Main

#include<bits/stdc++.h>
#define ll long long
#define pii pair<int,int>
#define pll pair<ll,ll>
#define plx pair<ll,int>
#define f first
#define s second
#define pb push_back
#define all(x) x.begin(),x.end()
#define vi vector<int>
#define vvi vector<vi>
#define pp pair<ll,int>
#define ub(x,i) upper_bound(all(x),i)-x.begin()
#define lb(x,i) lower_bound(all(x),i)-x.begin()
#define t3 tuple<int,int,int>
using namespace std;
const int mxn=1e5+5;
vector<pair<int,ll>>g[mxn];
int rt=-1;
ll ans=0;
ll rs1=-1,rs2=-1;
ll dp[mxn][2]{0};
ll add[mxn]{0};
int hv[mxn]{0},whv[mxn]{0};
int nd[mxn]{0};
int tin[mxn],tout[mxn],cur=0,sz[mxn]{0};
vector<ll>ord,ord2;
struct fenwick{
    ll fw[2*mxn]{0};
    void add(int i,ll amt){
        for(;i<2*mxn;i+=i&-i)fw[i]+=amt;
    }
    ll qr(int i,ll rs=0){
        for(;i;i-=i&-i)rs+=fw[i];
        return rs;
    }
}fen[2];
void dfs(int u,int p){
    dp[u][0]=dp[u][1]=0;
    for(auto [v,w]:g[u]){
        if(v==p)continue;
        dfs(v,u);sz[u]+=sz[v];
        ll t=dp[v][0]+w;
        if(t>dp[u][0])swap(t,dp[u][0]);
        if(t>dp[u][1])swap(t,dp[u][1]);
    }
    if(dp[u][0]+dp[u][1]>ans){
        ans=dp[u][0]+dp[u][1];
        rs1=dp[u][0];rs2=dp[u][1];
        rt=u;
    }
}bool vis[mxn]{0};
vector<pair<ll,int>>vec[mxn];
void dfs2(int u,int p,ll w){
    if(g[u].size()==1)vec[u].pb({0,u});
    tin[u]=++cur;sz[u]=1;hv[u]=-1;nd[tin[u]]=u;
    for(auto [v,w]:g[u]){
        if(v==p)continue;
        dfs2(v,u,w);sz[u]+=sz[v];
        if(hv[u]==-1||sz[hv[u]]<sz[v])hv[u]=v,whv[u]=w;
        if(vec[u].size()<vec[v].size())swap(vec[u],vec[v]);
        if(!vec[u].empty()&&!vec[v].empty()&&vec[u].back()>vec[v].back())vec[v].pb(vec[u].back()),vec[u].pop_back();
        for(auto it : vec[v])vec[u].pb(it);
    }vec[u].back().f+=w;ord.pb(vec[u].back().f);tout[u]=cur;
}
bool ch(int u,int p){
    for(auto [v,w]:g[u]){
        if(v==p)continue;
        if(ch(v,u))vis[u]=1;
    }return vis[u];
}ll res[mxn]{0};ll qs[mxn]{0};
vector<pair<ll,int>>vv[mxn];
ll sm[mxn]{0};
bool dele[mxn]{0};
vector<ll>ba[mxn];
void sol(int u,int p,int hd,ll d,bool del){
    if(g[u].size()==1)vv[u].pb({0,u});
    for(auto [v,w]:g[u]){
        if(v==p||!vis[v]||v==hv[u])continue;
        sol(v,u,v,w,1);
        if(vv[u].size()<vv[v].size())swap(vv[u],vv[v]);
        if(!vv[u].empty()&&!vv[v].empty()&&vv[u].back()>vv[v].back())vv[v].pb(vv[u].back()),vv[u].pop_back();
        for(auto it : vv[v])vv[u].pb(it);
        sm[u]+=sm[v];
    }
    if(hv[u]!=-1&&vis[hv[u]])sol(hv[u],u,hd,whv[u],0);
    if(hv[u]!=-1&&vis[hv[u]]){
        if(vv[u].size()<vv[hv[u]].size())swap(vv[u],vv[hv[u]]);
        if(!vv[u].empty()&&!vv[hv[u]].empty()&&vv[u].back()>vv[hv[u]].back())vv[hv[u]].pb(vv[u].back()),vv[u].pop_back();
        for(auto it : vv[hv[u]])vv[u].pb(it);
        sm[u]+=sm[hv[u]];
    }
    res[u]=ans-sm[u];int k=vv[u].size();
    for(auto [v,w]:g[u]){
        if(v==p||!vis[v]||v==hv[u])continue;
        for(int i=tin[v];i<=tout[v];i++)if(g[nd[i]].size()==1&&vis[nd[i]]){fen[0].add(ub(ord,add[nd[i]]),1),fen[1].add(ub(ord,add[nd[i]]),add[nd[i]]),ba[hd].pb(nd[i]);}
    }
    ll l=0,r=1e16;
    while(l<r){
        ll m=(l+r+1)>>1;
        int x=fen[0].qr(ub(ord,1e16))-fen[0].qr(lb(ord,m));
        x+=(int)ord2.size()-(lb(ord2,m));
        if(x>=k)l=m;
        else r=m-1;
    }res[u]+=qs[ord2.size()]-qs[lb(ord2,l)]+fen[1].qr(ub(ord,1e16))-fen[1].qr(lb(ord,l));
    ll x=(int)ord2.size() -(lb(ord2,l))+fen[0].qr(ub(ord,1e16))-fen[0].qr(lb(ord,l));
    res[u]-=(x-k)*l;
    if(g[u].size()!=1)fen[0].add(ub(ord,add[vv[u].back().s]),-1);
    if(g[u].size()!=1)fen[1].add(ub(ord,add[vv[u].back().s]),-add[vv[u].back().s]);
    sm[u]+=d;vv[u].back().f+=d;add[vv[u].back().s]+=d;
    fen[0].add(ub(ord,add[vv[u].back().s]),1);
    fen[1].add(ub(ord,add[vv[u].back().s]),add[vv[u].back().s]);
    ba[hd].pb(vv[u].back().s);
    if(del){
        for(auto it : ba[hd]){
            if(dele[it])continue;
            fen[0].add(ub(ord,add[it]),-1);
            fen[1].add(ub(ord,add[it]),-add[it]);
            dele[it]=1;
        }
        for(auto it : ba[hd]){
            dele[it]=0;
        }
    }
}
ll d[mxn]{0};
ll di[mxn]{0};
ll curmx=0,curme;
void cal(int u,int p){
    if(d[u]>curmx){
        curmx=d[u],curme=u;
    }
    for(auto [v,w]:g[u]){
        if(v==p)continue;
        d[v]=d[u]+w;
        cal(v,u);
    }
}
void solve1(int n){
    cal(1,1);
    int rrt=curme;
    memset(d,0,sizeof d);
    curmx=curme=0;
    cal(rrt,rrt);
    for(int i=1;i<=n;i++)di[i]=max(di[i],d[i]);
    memset(d,0,sizeof d);rrt=curme;
    cal(rrt,rrt);
    for(int i=1;i<=n;i++)di[i]=max(di[i],d[i]);
    for(int i=1;i<=n;i++)cout<<di[i]<<'\n';
    return;
}
void solve2(int u,int p,ll d){
    if(!vis[u])res[u]=res[p]+d;
    for(auto [v,w]:g[u]){
        if(v==p)continue;
        solve2(v,u,w);
    }
}
int main(){
    ios_base::sync_with_stdio(0);cin.tie(0);
    int n,k;cin>>n>>k;
    for(int i=1;i<=n-1;i++){
        int u,v,w;cin>>u>>v>>w;
        g[u].pb({v,w});
        g[v].pb({u,w});
    }if(k==1){solve1(n);return 0;}ord.pb(0);
    for(int i=1;i<=n;i++)if(g[i].size()>1)rt=i;
    dfs(rt,rt);dfs2(rt,rt,0);sort(all(ord));ord.erase(unique(all(ord)),ord.end());
    for(int j=0;j<vec[rt].size();j++){
        if(vec[rt][j].f==rs1){vis[vec[rt][j].s]=1,swap(vec[rt][j],vec[rt].back());break;}
    }vec[rt].pop_back();k-=2;
    for(int j=0;j<vec[rt].size();j++){
        if(vec[rt][j].f==rs2){vis[vec[rt][j].s]=1,swap(vec[rt][j],vec[rt].back());break;}
    }vec[rt].pop_back();sort(all(vec[rt]));
    while(k>0&&!vec[rt].empty()){
        ans+=vec[rt].back().f;
        vis[vec[rt].back().s]=1;vec[rt].pop_back();k--;
    }for(auto it : vec[rt])ord2.pb(it.f);
    for(int i=0;i<ord2.size();i++)qs[i+1]=qs[i]+ord2[i];
    ch(rt,rt);sol(rt,rt,rt,0,0);solve2(rt,rt,0);
    for(int i=1;i<=n;i++)cout<<res[i]<<'\n';
}

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32