BZOJ1468: Tree

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int MAXN=1e6+10;
const int INF=1e7+10;
inline char nc()
{   
    static char buf[MAXN],*p1=buf,*p2=buf;
    return p1==p2&&(p2=(p1=buf)+fread(buf,1,MAXN,stdin),p1==p2)?EOF:*p1++;
}
inline int read()
{
    char c=nc();int x=0,f=1;
    while(c<'0'||c>'9'){if(c=='-')f=-1;c=nc();}
    while(c>='0'&&c<='9'){x=x*10+c-'0',c=nc();}
    return x*f;
}
struct node
{
    int u,v,w,nxt; 
}edge[MAXN];
int head[MAXN];
int num=1;
inline void AddEdge(int x,int y,int z)
{
    edge[num].u=x;
    edge[num].v=y;
    edge[num].w=z;
    edge[num].nxt=head[x];
    head[x]=num++;
}
int K,F[MAXN],siz[MAXN],root,sum,vis[MAXN],tot[MAXN],cnt,deep[MAXN],ans=0;
void GetRoot(int now,int fa)
{
    siz[now]=1;
    for(int i=head[now];i!=-1;i=edge[i].nxt)
    {
        if(edge[i].v==fa||vis[edge[i].v]) continue;
        GetRoot(edge[i].v,now);
        siz[now]+=siz[edge[i].v];
        F[now]=max(F[now],siz[edge[i].v]);
    }
    F[now]=max(F[now],sum-siz[now]);
    if(F[now]<F[root]) root=now;
}
void GetDeep(int now,int fa)
{
    tot[++cnt]=deep[now];
    for(int i=head[now];i!=-1;i=edge[i].nxt)
    {
        if(vis[edge[i].v]||edge[i].v==fa) continue;
        deep[edge[i].v]=deep[now]+edge[i].w;
        GetDeep(edge[i].v,now);
    }
}
int Calc(int now,int val)//now点满足条件的个数 
{
    cnt=0;deep[now]=val;//一个小技巧 
    GetDeep(now,0);
    sort(tot+1,tot+cnt+1);
    int l=1,r=cnt,NowAns=0;
    while(l<r)
    {
        if(tot[l]+tot[r]<=K) NowAns+=r-l,l++;
        else r--;
    }
    return NowAns;
}
void Solve(int now)
{
    vis[now]=1;//别忘了打标记 
    ans+=Calc(now,0);
    
    for(int i=head[now];i!=-1;i=edge[i].nxt)
    {
        if(vis[edge[i].v]) continue;
        ans-=Calc(edge[i].v,edge[i].w);
        sum=siz[edge[i].v];
        root=0;
        GetRoot(edge[i].v,0);
        Solve(edge[i].v);
    }
}
int main()
{
    #ifdef WIN32
    freopen("a.in","r",stdin);
    #else
    #endif
    memset(head,-1,sizeof(head));
    int N=read();
    for(int i=1;i<=N-1;i++)
    {
        int x=read(),y=read(),z=read();
        AddEdge(x,y,z);
        AddEdge(y,x,z);
    }
    K=read();
    F[0]=INF;sum=N;
    //这里有个技巧，把root设置为0，f[0]=INF，那么可以解决找重心时的边界问题 
    GetRoot(1,0);
    Solve(root);
    printf("%d",ans);
    return 0;
}