洛谷 P4178 Tree

#include<iostream>
#include<cstdlib>
#include<cstdio>
#include<cmath>
#include<algorithm>
#include<cstring>
#define rg register
#define _ 40010
using namespace std;
int n,record[_],num_of_edges,size[_],dfn[_],dfns,F[_],root,dis[_];//size[]不多说,F[]是i点儿子中最大子树大小
bool passed[_];
int stack[_],top,ans,k;
struct pp
{
    int next,to,w;
}edge[_<<1];
inline int read()
{
    rg int save=0,w=1;rg char q=getchar();
    while(q<'0'||q>'9'){if(q=='-')w=-1;q=getchar();}
    while(q>='0'&&q<='9')save=(save<<3)+(save<<1)+q-'0',q=getchar();
    return save*w;
}
void getroot(rg int i,rg int fa,rg int s)
{
    size[i]=1,F[i]=0,dfn[i]=++dfns;
    for(rg int j=record[i];j;j=edge[j].next)
    {
        rg int to=edge[j].to;
        if(to!=fa&&!passed[to])
        {
            getroot(to,i,s);
            F[i]=max(F[i],size[to]);
            size[i]+=size[to];
        }
    }
    F[i]=max(F[i],s-size[i]);
    root=F[i]<F[root]?i:root;
}
void comput(rg int i,rg int dad)
{
    stack[++top]=dis[i];//stack直接把边权记下来,无需记点
    for(rg int j=record[i];j;j=edge[j].next)
    {
        rg int to=edge[j].to;
        if((!passed[to])&&dad!=to)
        {
            dis[to]=dis[i]+edge[j].w;
            comput(to,i);
        }
    }
}
inline int doit(rg int i,rg int d)//数出 该联通块中若使所有路径强制经过重心,<=k的有多少    (d为初始距离(去重用))
{
    top=0;dis[i]=d;
    comput(i,0);//计算各点到该重心的路径长
    sort(stack+1,stack+top+1);
    rg int l=1,r=top,result=0;
    while(l<r)
        if(stack[l]+stack[r]<=k)result+=r-l,l++;
        else r--;
    return result;
}
void dfs(rg int i,rg int all)
{
    ans+=doit(i,0);
    passed[i]=1;
    for(rg int j=record[i];j;j=edge[j].next)
    {
        rg int to=edge[j].to;
        if(!passed[to])
        {
            rg int s;
/* ☆☆☆*/    ans-=doit(to,edge[j].w);//去掉那些 在to的子树加上了[i->to的边权]被算作i的贡献的路径(这些路径本不存在)
            if(dfn[i]<dfn[to])s=size[to];
            else s=all-size[i];
            root=0;dfns=0;
            getroot(to,i,s);
//            printf("%d\n",root);
            dfs(to,s);
        }
    }
}
inline void add(rg int from,rg int to,rg int ww)
{
    edge[++num_of_edges]=(pp){record[from],to,ww};
    record[from]=num_of_edges;
}
int main()
{
    n=read();
//    cout<<"___\n"<<n<<endl;
    rg int i,j;
    for(i=1;i<n;++i)
    {
        rg int x=read(),y=read(),ww=read();
        add(x,y,ww),add(y,x,ww);
//        cout<<x<<' '<<y<<' '<<ww<<endl;
    }
    k=read();
//    cout<<k<<endl;
    root=0;F[0]=n;//操作 "root=F[i]<F[root]?i:root" 需要
    getroot(1,0,n);
    dfs(root,n);
    printf("%d\n",ans);
    return 0;
}