poj1741 树分治

题目链接：http://poj.org/problem?id=1741

可参考论文：https://wenku.baidu.com/view/e087065f804d2b160b4ec0b5.html

题意：求树上的点对（u,v）个数，点对满足dist(u,v)<=k。其中，dist(u,v)表示u到v的距离

解法：树上点分治。

树上递归求解问题，为了防止递归过程中树退化成一条链，我们每次递归时找树的重心。重心求法可以参考这篇博客http://blog.csdn.net/zt_1995/article/details/58729558。

这样可以保证递归的深度为log(n)。递归过程中，我们选子树的重心作为根节点，那么将会有两种点对：

1.点对连线不经过根，那么点对的最近公共祖先是根的儿子节点

2.点对连线经过根

（1）对于这个根节点，我们只需要求出点对连线经过根的情况数（即情况1），因为点对连线不经过根的情况数（情况2）在递归遍历子树的时候会处理。

（2）可以求出子树上所有点到根的距离，排序，然后o（n）的复杂度求出所有满足dist(u,v)<=k的情况数。

根节点的情况数（1）=根节点的点对（2）-根所有儿子的点对（2）

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<queue>
#include<stdlib.h>
#define inf 100000000
using namespace std;

const int N=10005;

int root;//重心
int siz;//子树大小
int son[N];
int dis[N],num;
int ans=0;
int head[N],top=0;
int n,k;
bool vis[N];
int far[N];
int size[N],maxv[N];
int f[N];
int que[N];
struct Edge
{
    int v,next;
    int w;
}edge[N*2];

void init()
{
    memset(head,-1,sizeof(head));
    top=0;
    siz=inf;
    ans=0;
    memset(vis,0,sizeof(vis));
}

void addedge(int u,int v,int w)
{
    edge[top].v=v;
    edge[top].next=head[u];
    edge[top].w=w;
    head[u]=top++;
}
//两种方法找重心
//dfs找重心
void getroot(int u,int fa)
{
    son[u]=1;f[u]=0;
//    int Max=0;
    for(int i=head[u] ; i!=-1 ; i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==fa||vis[v]) continue;
        getroot(v,u);
        son[u]+=son[v];
        f[u]=max(f[u],son[v]);
    }
    f[u]=max(f[u],n-son[u]);
    if(f[u]<f[root]){
        root=u;
    }
}
//bfs找重心
void getroot(int u)
{
    int tail=0;
    que[tail++]=u;
    far[u]=0;
    siz=inf;
    root=u;
    for(int k=0;k<tail;k++){
        int x = que[k];
        for(int i=head[x] ; i!=-1 ; i=edge[i].next){
            int v=edge[i].v;
            if(vis[v]||v==far[x])continue;
            far[v]=x;
            que[tail++]=v;
        }
    }

    for(int k=tail-1;k>=0;k--){
        int x = que[k];
        son[x]=1;
        int Max=0;
        for(int i=head[x];i!=-1;i=edge[i].next){
            int v=edge[i].v;
            if(vis[v]||v==far[x])continue;
            son[x]+=son[v];
            Max=max(Max,son[v]);
        }
        Max=max(Max,n-son[x]);
        if(Max<siz) siz=Max,root=x;
    }
}

void getdis(int u,int fa,int d)
{
    dis[++num]=d;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==fa||vis[v])continue;
        getdis(v,u,d+edge[i].w);
    }
}

int cal(int u,int d)
{
    num=0;
    getdis(u,u,d);
    int cnt=0;
    sort(dis+1,dis+num+1);
    int i=1,j=num;
    while(i<j)
    {
        while(dis[i]+dis[j]>k&&i<j){
            j--;
        }
        cnt+=j-i;
        i++;
    }
    return cnt;
}

void work(int u)
{
    getroot(u,root=0);
    ans+=cal(root,0);
    vis[root]=1;

    for(int i=head[root];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(vis[v]) continue;
        ans-=cal(v,edge[i].w);
        f[0]=n=son[v]; //更新子树大小。。非常重要，因为这个找了一天bug
        work(v);
    }
}

int main()
{
    freopen("in.cpp","r",stdin);
    while(~scanf("%d%d",&n,&k))
    {
        if(!(n||k))break;
        init();
        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);
        }
        f[0]=n;
        work(1);
        printf("%d\n",ans);
    }
    return 0;
}