计蒜客模拟赛 #5 (B 题) 动态点分治+线段树

虽然是裸的换根dp，但是为了在联赛前锻炼码力，强行上了点分树+线段树.   

写完+调完总共花了不到 $50$ 分钟，感觉还行.  

code: 

#include <bits/stdc++.h>   
#define N 420004  
#define LL long long  
#define setIO(s) freopen(s".in","r",stdin) 
using namespace std;  
namespace IO 
{
    char *p1,*p2,buf[100000];
    #define nc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,100000,stdin),p1==p2)?EOF:*p1++)
    int rd() {int x=0; char c=nc(); while(c<48) c=nc(); while(c>47) x=(((x<<2)+x)<<1)+(c^48),c=nc(); return x;}
};  
int n,K,edges; 
int hd[N],nex[N<<1],val[N<<1],to[N<<1],F[N],G[N];      
inline void add(int u,int v,int c) 
{
    nex[++edges]=hd[u],hd[u]=edges,to[edges]=v,val[edges]=c;    
}   
struct tree 
{  
    LL dis[N];  
    int size[N],son[N],top[N],fa[N],dep[N];     
    void dfs1(int u,int ff) 
    { 
        fa[u]=ff,dep[u]=dep[ff]+1,size[u]=1; 
        for(int i=hd[u];i;i=nex[i]) 
        {
            int v=to[i]; 
            if(v==ff) continue;
            dis[v]=dis[u]+val[i];       
            dfs1(v,u);       
            size[u]+=size[v];   
            if(size[v]>size[son[u]])   son[u]=v; 
        }          
    }   
    void dfs2(int u,int tp) 
    {
        top[u]=tp;    
        if(son[u])  dfs2(son[u],tp); 
        for(int i=hd[u];i;i=nex[i]) 
            if(to[i]!=fa[u]&&to[i]!=son[u])    dfs2(to[i],to[i]);  
    }  
    inline int LCA(int x,int y) 
    {
        while(top[x]!=top[y]) 
        {
            dep[top[x]]>dep[top[y]]?x=fa[top[x]]:y=fa[top[y]];  
        } 
        return dep[x]<dep[y]?x:y;  
    } 
    inline LL Dis1(int x,int y) 
    {
        return dis[x]+dis[y]-(dis[LCA(x,y)]<<1);  
    } 
    inline LL Dis2(int x,int y) 
    {
        return dep[x]+dep[y]-(dep[LCA(x,y)]<<1);   
    }
}tree;  
struct seg
{    
    int tot; 
    inline void init() { tot=0;}    
    inline int newnode() { return ++tot; }   
    struct node
    {   
        int ls,rs,s2; 
        LL sum; 
    }t[N*20];     
    void update(int &x,int l,int r,int p,LL v) 
    {
        if(!x) x=newnode();     
        t[x].sum+=v;         
        t[x].s2++;  
        if(l==r) return; 
        int mid=(l+r)>>1;            
        if(p<=mid)    update(t[x].ls,l,mid,p,v);   
        else     update(t[x].rs,mid+1,r,p,v);  
    }   
    LL query(int x,int l,int r,int L,int R) 
    {     
        if(!x) return 0;    
        if(l>=L&&r<=R)    return t[x].sum;  
        int mid=(l+r)>>1;  
        LL re=0ll; 
        if(L<=mid)    re+=query(t[x].ls,l,mid,L,R);                   
        if(R>mid)     re+=query(t[x].rs,mid+1,r,L,R);  
        return re;   
    }
    int q2(int x,int l,int r,int L,int R) 
    {
        if(!x) return 0; 
        if(l>=L&&r<=R)    return t[x].s2;    
        int mid=(l+r)>>1,re=0; 
        if(L<=mid)   re+=q2(t[x].ls,l,mid,L,R);  
        if(R>mid)    re+=q2(t[x].rs,mid+1,r,L,R);  
        return re;   
    }
}seg; 
struct opt 
{   
    int root,sn; 
    int Fa[N],size[N],vis[N],mx[N];    
    void getroot(int u,int ff) 
    {     
        size[u]=1,mx[u]=0;   
        for(int i=hd[u];i;i=nex[i]) 
        {
            int v=to[i]; 
            if(v==ff||vis[v]) continue;  
            getroot(v,u);  
            size[u]+=size[v];            
            mx[u]=max(mx[u], size[v]); 
        }  
        mx[u]=max(mx[u], sn-size[u]);   
        if(mx[u]<mx[root])  root=u;    
    }         
    void calc(int pp,int u,int lst,int ff,int dep,LL w) 
    {          
        if(dep<=K) seg.update(F[pp],0,K,dep,w);    
        if(ff) 
        {
            LL tmp2=tree.Dis1(u,ff);    
            int tmp1=tree.Dis2(u,ff);              
            if(tmp1<=K)     seg.update(G[pp],0,K,tmp1,tmp2);        
        }   
        for(int i=hd[u];i;i=nex[i])   
        {
            int v=to[i];   
            if(vis[v]||v==lst) continue;  
            calc(pp,v,u,ff,dep+1,w+1ll*val[i]);                         
        }
    }  
    void dfs(int u) 
    {    
        vis[u]=1;    
        calc(u,u,0,Fa[u],0,0ll);    
        for(int i=hd[u];i;i=nex[i]) 
        {
            int v=to[i];      
            if(vis[v]) continue;    
            root=0,sn=size[v],getroot(v,u),Fa[root]=u,dfs(root);       
        }
    } 
    LL getdis(int u) 
    {
        LL tmp=seg.query(F[u],0,K,0,K);                             
        int U=u;       
        for(;Fa[u];u=Fa[u])        
        {
            int tt=tree.Dis2(U,Fa[u]);                
            if(tt<=K) 
            {  
                LL a=seg.query(F[Fa[u]],0,K,0,K-tt)-seg.query(G[u],0,K,0,K-tt);  
                LL b=seg.q2(F[Fa[u]],0,K,0,K-tt)-seg.q2(G[u],0,K,0,K-tt);  
                tmp+=a+b*tree.Dis1(U,Fa[u]);          
            }
        }    
        return tmp;   
    }
    void solve() 
    {
        root=0;  
        mx[0]=sn=n;           
        getroot(1,0);       
        dfs(root);        
        int i,j;           
        for(i=1;i<=n;++i) 
        {
            // getdis(i); 
            printf("%lld ",getdis(i));  
        }       
    }
}opt;  
int main() 
{  
    //  setIO("input"); 
    using namespace IO;  
    int i,j;  
    n=rd(),K=rd(); 
    //   scanf("%d%d",&n,&K);   
    for(i=1;i<n;++i) 
    { 
        int x,y,c; 
        x=rd(),y=rd(),c=rd(); 
        add(x,y,c),add(y,x,c); 
    }    
    tree.dfs1(1,0); 
    tree.dfs2(1,1);     
    seg.init();           
    opt.solve();     
    return 0;  
}