POJ 1741.Tree 树分治 树形dp 树上点对

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 24258		Accepted: 8062

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

LouTiancheng@POJ

 

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

题意：一棵有n个结点的树，求距离不超过k的点对

思路：树分治。n比较大，直接枚举所有点对肯定是不行的。按照重心把树分成若干子树，那么所有的点对一定属于 1）点u、v属于同一子树的点对； 2）点u、v属于不同子树的定点对； 3）重心s和其他点组成点对。1）情况可以通过递归得到。2）情况，只要先求出每个点到重心s的距离，就可以统计出和不超过k的点对数。而3）情况，添加一个0的顶点，就成为了情况2）。需要注意的是，需要避免重复统计，即应该在1）中统计的属于同一子树的点对，要避免在2）中进行统计。递归深度最多为log(n)层，每层总共有n个结点。

代码：

#include<iostream>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
#include<set>
#include<bitset>
#include<map>
#include<queue>
#include<stack>
#include<vector>
using namespace std;
#define bug(x) cout<<"bug"<<x<<endl;
#define PI acos(-1.0)
#define eps 1e-8
typedef long long ll;
typedef pair<int,int> P;
const int N=1e5+100,M=1e5+100;
const int inf=0x3f3f3f3f;
const ll INF=1e18+7,mod=1e9+7;
struct edge
{
    int from,to;
    int w;
    int next;
};
edge es[M];
int cut,head[N];
int si[N],maxx[N];
bool vis[N];
int deep[N];
int k;
int root,ans;
void init()
{
    cut=0;
    memset(head,-1,sizeof(head));
}
void addedge(int u,int v,int w)
{
    cut++;
    es[cut].from=u,es[cut].to=v;
    es[cut].w=w;
    es[cut].next=head[u];
    head[u]=cut;
}
int getroot(int u,int fa,int n)
{
    si[u]=1,maxx[u]=0;
    for(int i=head[u]; i!=-1; i=es[i].next)
    {
        int v=es[i].to;
        if(v==fa||vis[v]) continue;
        si[u]+=getroot(v,u,n);
        maxx[u]=max(maxx[u],si[v]);
    }
    maxx[u]=max(maxx[u],n-si[u]);
    if(maxx[u]<maxx[root]) root=u;
    return si[u];
}
void getdeep(int u,int fa,int d)
{
    deep[++deep[0]]=d;
    for(int i=head[u]; i!=-1; i=es[i].next)
    {
        edge e=es[i];
        if(e.to==fa||vis[e.to]) continue;
        getdeep(e.to,u,d+e.w);
    }
}
int cal(int u,int fa,int d)
{
    deep[0]=0;
    getdeep(u,fa,d);
    sort(deep+1,deep+deep[0]+1);
    int l=1,r=deep[0];
    int res=0;
    while(l<r)
    {
        if(deep[l]+deep[r]<=k) res+=r-l,l++;
        else r--;
    }
    return res;
}
void solve(int u)
{
    vis[u]=true;
    ans+=cal(u,0,0);///统计符合情况的点对数
    for(int i=head[u]; i!=-1; i=es[i].next)
    {
        edge e=es[i];
        if(vis[e.to]) continue;
        ans-=cal(e.to,0,e.w);///删除同一子树的点对数
        root=0;
        getroot(e.to,0,si[e.to]);
        solve(root);///递归同一子树
    }
}
int main()
{
    int n;
    while(scanf("%d%d",&n,&k)!=EOF)
    {
        if(n==0&&k==0) break;
        init();
        for(int i=1; i<n; i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            addedge(u,v,w);
            addedge(v,u,w);
        }
        memset(vis,false,sizeof(vis));
        root=0,maxx[0]=inf;
        ans=0;
        getroot(1,0,n);
        solve(root);
        printf("%d
",ans);
    }
    return 0;
}