【树上点分治模板】

POJ1741

给定一棵树（10000个点），求其u，v两点距离大于等于k的对数。

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
typedef long long LL;
const int INF=0x3f3f3f3f;
const int N=10010;
int root,maxx,S,cnt,n,k;
int sum[N],maxs[N],d[N];
LL ans;
bool vis[N];
struct edge
{
    int from,to,cost,next;
}e[N*2];
int head[N],nu;
void add(int x,int y,int val)
{
    e[nu].to=y;
    e[nu].cost=val;
    e[nu].next=head[x];
    head[x]=nu++;
}
void init()
{
    memset(vis,false,sizeof(vis));
    memset(head,-1,sizeof(head));
    nu=0;
}
void getroot(int x,int fa) //寻找树的重心过程
{
    sum[x]=1;
    maxs[x]=0;
    for(int i=head[x];i+1;i=e[i].next)
    {
        int y=e[i].to;
        if(y==fa||vis[y])continue;
        getroot(y,x);
        sum[x]+=sum[y];
        maxs[x]=max(maxs[x],sum[y]);
    }
    maxs[x]=max(maxs[x],S-sum[x]);
    if(maxs[x]<maxx)
    {
        maxx=maxs[x];
        root=x;
    }
}
void dfs(int x,int fa,int len)  //对于每个点枚举其儿子，以达到枚举路径的目的。
{
    d[++cnt]=len;
    for(int i=head[x];i+1;i=e[i].next)
    {
        int y=e[i].to;
        if(vis[y]||y==fa)continue;
        dfs(y,x,len+e[i].cost);
    }
}
int solve(int x,int val)
{
    cnt=0;
    dfs(x,-1,val);
    sort(d+1,d+cnt+1);
    int ret=0,l=1,r=cnt;
    while(l<r)
    {
        if(d[l]+d[r]<=k)
        {
            ret+=r-l;
            l++;
        }
        else
        r--;
    }
    return ret;
}
void divtree(int x,int fa)
{
    ans=ans+solve(x,0);
    vis[x]=true;
    for(int i=head[x];i+1;i=e[i].next)
    {
        int y=e[i].to;
        if(y==fa||vis[y])continue;
        ans=ans-solve(y,e[i].cost);
        maxx=INF;
        S=sum[y];
        root=0;
        getroot(y,-1);
        divtree(root,-1);
    }
}
int main()
{
    int x,y,z;
    while(scanf("%d%d",&n,&k)!=EOF&&(n||k))
    {
        init();
        for(int i=1;i<n;i++)
        {
            scanf("%d%d%d",&x,&y,&z);
            add(x,y,z);
            add(y,x,z);
        }
        maxx=INF;
        S=n;
        root=0;
        ans=0;
        getroot(1,0);
        divtree(root,0);
        printf("%lld\n",ans);
    }
}