POJ 1741 Tree 点分治

http://poj.org/problem?id=1741

题目大意：给一棵树，问有多少对点对$(x,y)$满足树上从$x$到$y$的距离$<=k$。$(x,y)$和$(y,x)$认为是同一对对点。
思路：点分治模板题 。

#include<cstdio>
#include<cstring>
#include<algorithm>
#define INF 0x3f3f3f3f
using namespace std;
typedef long long ll;

const int maxn=1e4+5;

struct Edge
{
    int to,nxt,dis;
}edge[maxn<<1];

ll re; //记录答案
int n,k,tot;
int head[maxn],dis[maxn],siz[maxn],vis[maxn],b[maxn];

inline void addedge(int u,int v,int d)
{
    edge[++tot].to=v,edge[tot].nxt=head[u],edge[tot].dis=d,head[u]=tot;
}

struct Tree
{
    int n,rt,num;
    void dfs(int u,int fa)
    {
        siz[u]=1;
        int MAX=0;
        for(int i=head[u];i;i=edge[i].nxt)
        {
            int v=edge[i].to;
            if(v==fa||vis[v])
                continue;
            dfs(v,u);
            siz[u]+=siz[v];
            MAX=max(MAX,siz[v]);
        }
        MAX=max(MAX,n-siz[u]);
        if(MAX<num)
            num=MAX,rt=u;
    }
    int getrt(int x)
    {
        rt=0,num=INF;
        dfs(x,0);
        return rt;
    }
}tree;

void getdis(int u,int fa)
{
    b[++b[0]]=dis[u];
    for(int i=head[u];i;i=edge[i].nxt)
    {
        int v=edge[i].to;
        if(v==fa||vis[v])
            continue;
        dis[v]=dis[u]+edge[i].dis;
        getdis(v,u);
    }
}

ll cal(int u,int val)
{
    dis[u]=val,b[0]=0;
    getdis(u,0);
    sort(b+1,b+1+b[0]);
    ll ans=0;
    int l=1,r=b[0];
    while(l<r)
    {
        if(b[l]+b[r]<=k)
            ans+=r-l,++l;
        else
            --r;
    }
    return ans;
}

void work(int u)
{
    vis[u]=1;
    re+=cal(u,0);
    for(int i=head[u];i;i=edge[i].nxt)
    {
        int v=edge[i].to;
        if(vis[v])
            continue;
        re-=cal(v,edge[i].dis);
        tree.n=siz[v];
        work(tree.getrt(v));
    }
}

int main()
{
    while(~scanf("%d %d",&n,&k)&&(n+k))
    {
        tot=0;
        memset(head,0,sizeof(head));
        int u,v,w;
        for(int i=1;i<n;i++)
        {
            scanf("%d %d %d",&u,&v,&w);
            addedge(u,v,w),addedge(v,u,w);
        }
        re=0,tree.n=n;
        memset(vis,0,sizeof(vis));
        work(tree.getrt(1));
        printf("%lld\n",re);
    }
    return 0;
}