POJ 1741 Tree

Description
Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.

---

【题目分析】
树的点分治。抄来做模板，慢慢理解。

---

【代码】

#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;
int n,k,e,ans,root,num;
int to[50001],ne[50001],w[50001],h[50001],vis[50001];
inline void add2(int a,int b,int c)
{
    to[e]=b;ne[e]=h[a];w[e]=c;h[a]=e++;
    to[e]=a;ne[e]=h[b];w[e]=c;h[b]=e++;
}
int mx[50001],size[50001],mi,dis[50001];
void dfssize(int u,int fa)
{
    size[u]=1,mx[u]=0;
    for (int i=h[u];i>=0;i=ne[i])
    {
        int v=to[i];
        if (v!=fa&&!vis[v])
        {
            dfssize(v,u);
            size[u]+=size[v];
            if (size[v]>mx[u]) mx[u]=size[v];
        }
    }
}
void dfsroot(int r,int u,int fa)
{
    if(size[r]-size[u]>mx[u]) mx[u]=size[r]-size[u];
    if (mx[u]<mi) mi=mx[u],root=u;
    for (int i=h[u];i>=0;i=ne[i])
    {
        int v=to[i];
        if (v!=fa&&!vis[v]) dfsroot(r,v,u);
    }
}
void dfsdis(int u,int d,int fa)
{
    dis[num++]=d;
    for (int i=h[u];i>=0;i=ne[i])
    {
        int v=to[i];
        if (v!=fa&&!vis[v]) dfsdis(v,d+w[i],u);
    }
}
int calc(int u,int d)
{
    int ret=0;num=0;
    dfsdis(u,d,0);
    sort(dis,dis+num);
    int i=0,j=num-1;
    while (i<j)
    {
        while (dis[i]+dis[j]>k&&i<j) j--;
        ret+=j-i;
        i++;
    }
    return ret;
}
void dfs(int u)
{
    mi=n;
    dfssize(u,0),dfsroot(u,u,0);
    ans+=calc(root,0);
    vis[root]=1;
    for (int i=h[root];i>=0;i=ne[i])
    {
        int v=to[i];
        if (!vis[v])ans-=calc(v,w[i]),dfs(v);
    }
}
int main()
{
    while (scanf("%d%d",&n,&k)!=EOF&&n&&k)
    {
        memset(vis,0,sizeof vis);
        memset(h,-1,sizeof h);
        e=ans=0;
        int a,b,c;
        for (int i=1;i<n;++i)
            scanf("%d%d%d",&a,&b,&c),add2(a,b,c);
        dfs(1);
        printf("%d\n",ans);
    }
}