Tree（树分治入门）

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

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 35091		Accepted: 11718

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.

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.

Output

For each test case output the answer on a single line.

Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

Source

题意：一棵有n个节点的树，每条边有个权值代表相邻2个点的距离，要求求出所有距离不超过k的点对(u,v)

题解：树分治
假设树以root为根节点，那么满足要求的点对有2种情况：
①路径经过root且dis(u,v)<=k
②路径不经过root，即其路径的最高点为子树上某一节点

对于第②种情况可以通过递归求解，这里只讨论第一种情况
该如何求解路径经过root且dis(u,v)<=k的合法点对数呢？

设dir[u]为u到根节点root的距离，那么只有满足dir[u]+dir[v]<=k且LCA(u,v)==root的点对才是合法的，
设cnt1=树中所有dis(u,v)<=k的点对数，cnt2=LCA(u,v)==root的子节点的合法点对数
那么以root为根的树种合法点对数为：ans=cnt1-cnt2
找出有多少个dir[u]+dir[v]的方法很简单：只需要排序后扫一遍即可。

总结一下算法的过程：
①计算以u为根的树种每棵子树的大小
②根据子树大小找出树的重心root(以树的重心为根的树，可以使其根的子树中节点最多的子树的节点最少)
③以root为根，计算树中每个点到root的距离dir
④计算树中所有满足dir[u]+dir[v]<=k的点对数cnt1
⑤计算以root的子节点为根的子树中，满足dir[u]+dir[v]<=k的点对数cnt2
⑥ans+=cnt1-cnt2
注意：每次计算完cnt1后，要将vis[root]=1，这样就可以将一棵树分解成若干棵子树

看代码：

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
using namespace std;
typedef long long LL;
const int maxn=2e4+5;
const int INF=1e9+7;
int cnt=0;
int head[maxn<<1];
int sonnum[maxn<<1],sonmax[maxn<<1];
int mi,pos;
int N,K;
int sum=0;
bool vis[maxn<<1];
LL ans;
vector<int>dis;
struct Edge
{
    int next,to,w;
}e[maxn<<1];
void Init()
{
    ans=0;cnt=0;
    for(int i=0;i<maxn;i++)
    {
        head[i]=-1;sonnum[i]=sonmax[i]=0;
        vis[i]=false;
    }
}
void add_edge(int u,int v,int w)
{
    e[++cnt].to=v;
    e[cnt].w=w;
    e[cnt].next=head[u];
    head[u]=cnt;
}
void Query_size(int root,int pre)//求当前树的子树大小
{
    sum++;//存树的大小
    sonnum[root]=1;sonmax[root]=0;//注意初始化 
    for(int i=head[root];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(vis[v]||v==pre) continue;
        Query_size(v,root);
        sonnum[root]+=sonnum[v];//子树节点有多少个
        sonmax[root]=max(sonmax[root],sonnum[v]);//最大的子树节点个数
    }
}
void Query_root(int root,int pre,int sum)//求当前树的重心
{
    for(int i=head[root];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(vis[v]||v==pre) continue;
        Query_root(v,root,sum);
    }
    int ma=max(sonmax[root],sum-sonnum[root]);
    if(mi>ma)
    {
        mi=ma;pos=root;
    }
}
void Query_dis(int root,int pre,int d)
{
    dis.push_back(d);
    for(int i=head[root];i!=-1;i=e[i].next)
    {
        int v=e[i].to,w=e[i].w;
        if(vis[v]||v==pre) continue;
        Query_dis(v,root,d+w);
    }
}
int cal(int root,int d)
{
    int ret=0;
    dis.clear();//存所有子节点到本身的距离
    Query_dis(root,0,d);
    sort(dis.begin(),dis.end());
    int i=0,j=dis.size()-1;
    while(i<j)
    {
        while(i<j&&dis[i]+dis[j]>K) j--;
        ret+=j-i;
        i++;
    }
    return ret;
}
void dfs(int root,int pre)
{
    sum=0;
    mi=INF;
    Query_size(root,pre);
    Query_root(root,pre,sum);//
    int rt=pos;
    ans+=cal(rt,0);//pos为找到的重心
    vis[rt]=true;//一定要标记 否则会往回走
    for(int i=head[rt];i!=-1;i=e[i].next)
    {
        int v=e[i].to,w=e[i].w;
        if(vis[v]) continue;
        ans-=cal(v,w);
        dfs(v,rt);
    }
}
int main()
{
    while(scanf("%d%d",&N,&K)!=EOF)
    {
        Init();
        if(N==0&&K==0) break;
        for(int i=1;i<N;i++)
        {
            int u,v,w;scanf("%d%d%d",&u,&v,&w);
            add_edge(u,v,w);
            add_edge(v,u,w);
        }
        dfs(1,0);
        printf("%lld
",ans);
    }
    return 0;
}