BZOJ 3784: 树上的路径

/**************************************************************
    Problem: 3784
    User: BeiYu
    Language: C++
    Result: Accepted
    Time:4540 ms
    Memory:26660 kb
****************************************************************/
 
#include <bits/stdc++.h>
using namespace std;
 
#define debug(a) cout<<#a<<"="<<a<<" "
#define mpr make_pair
typedef pair< int,int > pr;
typedef pair< int,pr > prr;
const int N = 100500;
 
int n,k,rt,rtt;
int d[N],sz[N],ud[N];
int vis[N<<1];
 
int cnte;
struct Edge { int to,w; }edge[N<<2];
vector< int > g[N];
priority_queue< pr > qq;
priority_queue< pr > q[N];
vector< int > nxt[N];
vector< int > tmp[N][2];
 
inline int in(int x=0,char ch=getchar()) { while(ch>‘9‘ || ch<‘0‘) ch=getchar();
    while(ch>=‘0‘ && ch<=‘9‘) x=x*10+ch-‘0‘,ch=getchar();return x; }
 
void AddEdge(int fr,int to,int w) { g[fr].push_back(cnte);edge[cnte++]=(Edge){ to,w }; }
 
void Print() {
    cout<<"start"<<endl;
    for(int i=1;i<=n;i++) {
        cout<<i<<"-->";
        for(int j=0;j<(int)g[i].size();j++) cout<<"id:"<<g[i][j]<<" "<<edge[g[i][j]].to<<" ";
        cout<<endl;
    }
    cout<<"-------------------------"<<endl;
}
 
void ReBuild(int u,int fa) {
    int cnt=0,id=-1,lst=0,fst=0;
    if(g[u].size()>3) for(int i=g[u].size()-1,v;i>=1;i--) 
        if((v=edge[g[u][i]].to)!=fa) {
            if(cnt&1) {
                g[n].push_back(g[u][i]);
            }else {
                g[++n].push_back(g[u][i]);
                if(lst) {
                    AddEdge(lst,n,0);
                    AddEdge(n,lst,0);
                }else {
                    fst=n;
                }lst=n;
            }
            cnt^=1;
            edge[g[u][i]^1].to=n;
            g[u].pop_back();
        }else id=g[u][i],g[u].pop_back();
    if(id!=-1) g[u].push_back(id);
    if(fst) AddEdge(u,fst,0),AddEdge(fst,u,0);
    for(int i=0;i<(int)g[u].size();i++) if(edge[g[u][i]].to!=fa) ReBuild(edge[g[u][i]].to,u);
     
}
 
void GetSize(int u,int fa) {
    d[u]=d[fa]+1,sz[u]=1;
    for(int i=g[u].size()-1,v;~i;i--) if((v=edge[g[u][i]].to)!=fa && !vis[g[u][i]]) {
        GetSize(v,u),sz[u]+=sz[v];
    }
}
void GetRoot(int u,int nn) {
    rt=-1,rtt=N;GetSize(u,u);
//  for(int i=1;i<=n;i++) cout<<sz[i]<<" ";cout<<endl;
//  for(int i=1;i<=n;i++) cout<<d[i]<<" ";cout<<endl;
    queue< pr > q;q.push(mpr(u,u));
    for(pr x;!q.empty();) {
        x=q.front(),q.pop();
        int u=x.first,tmp;
        for(int i=0,v;i<(int)g[u].size();i++) if((v=edge[g[u][i]].to)!=x.second && !vis[g[u][i]]){
            if(d[u]>d[v]) tmp=max(sz[u],nn-sz[u]);
            else tmp=max(sz[v],nn-sz[v]);
//          debug(u),debug(v),debug(tmp),debug(g[u][i])<<endl;
            if(tmp<rtt) rtt=tmp,rt=g[u][i];
            q.push(mpr(v,u));
        }
    }
//  debug(u),debug(rt)<<endl;
}
void GetDep(int u,int fa,int w,vector< int > &vv) {
    vv.push_back(w);
    for(int i=0,v;i<(int)g[u].size();i++) if((v=edge[g[u][i]].to)!=fa && !vis[g[u][i]]) {
        GetDep(v,u,w+edge[g[u][i]].w,vv);
    }
}
 
void GetAns(int x,int nn) {
    vis[x]=vis[x^1]=1;
    int u=edge[x].to,v=edge[x^1].to;
    GetDep(u,u,0,tmp[x][0]),GetDep(v,v,0,tmp[x][1]);
     
    sort(tmp[x][0].begin(),tmp[x][0].end(),greater< int >());
    sort(tmp[x][1].begin(),tmp[x][1].end(),greater< int >());
     
//  cout<<"---------------------"<<endl;
//  debug(x),debug(nn),debug(u),debug(v)<<endl;
//  for(int i=0;i<(int)tmp[x][0].size();i++) cout<<tmp[x][0][i]<<" ";cout<<endl;
//  for(int i=0;i<(int)tmp[x][1].size();i++) cout<<tmp[x][1][i]<<" ";cout<<endl;
 
     
    nxt[x].resize(tmp[x][0].size());
    for(int i=0;i<(int)tmp[x][0].size();i++) nxt[x][i]=0;
     
    for(int i=0;i<(int)tmp[x][0].size();i++) q[x].push(mpr(tmp[x][0][i]+tmp[x][1][0]+edge[x].w,i));
     
//  debug(tmp[x][0][i]+tmp[x][1][0]+edge[x].w)
//  debug(q[x].top().first)<<endl;
    qq.push(mpr(q[x].top().first,x));
     
    if(d[u]<d[v]) swap(u,v);
    int ss=sz[u];
    GetRoot(u,ss);
    if(rt!=-1) GetAns(rt,ss);
    GetRoot(v,nn-ss);
    if(rt!=-1) GetAns(rt,nn-ss);
}
void init() {
    n=in(),k=in();
    for(int i=1,u,v,w;i<n;i++)
        u=in(),v=in(),w=in(),AddEdge(u,v,w),AddEdge(v,u,w);
     
//  ReBuild(1,1);
     
//  Print();
     
//  cout<<"qwq"<<endl;
     
    GetRoot(1,n);
     
//  debug(rt)<<endl;
     
    GetAns(rt,n);
     
}
void Del(int x) {
    pr r=q[x].top();
    q[x].pop();
    int u=r.second;
    if(nxt[x][u]+1<(int)tmp[x][1].size()) nxt[x][u]++,q[x].push(mpr(tmp[x][0][u]+tmp[x][1][nxt[x][u]]+edge[x].w,u));
    qq.push(mpr(q[x].top().first,x));
}
int Query() {
    pr x=qq.top();qq.pop();
    int r=x.first,u=x.second;
    Del(u);
     
    return r;
}
int main() {
//  freopen("in.in","r",stdin);
//  ios::sync_with_stdio(false);
    init();
     
    for(;k--;) {
        printf("%d\n",Query());
    }
     
    return 0;
}