计蒜客模拟赛 #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;
}