POJ 1741 Tree(点分治点对<=k)

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

题解

求树上总共有多少点对(u,v)距离<=k

题意

初学点分治

树的重心定义：找到一个点,其所有的子树中最大的子树节点数最少,那么这个点就是这棵树的重心,删去重心后，生成的多棵树尽可能平衡

点分治就是把树按子树重心不断DFS，然后处理每个子树对答案的贡献，会用到容斥，把n^2的复杂度优化为nlogn

处理每个子树对答案的贡献，先求出重心到各个点的距离，然后sort一下，然后可以二分出满足<=k的答案

代码

#include<stdio.h>
#include<string.h>
#include<vector>
#include<algorithm>
using namespace std;

const int maxn=1e4+5;

vector< pair<int,int> >G[maxn];
int mx[maxn],size[maxn],vis[maxn],dis[maxn],ans,MIN,n,k,num,root;
void dfssize(int u,int fa)
{
    size[u]=1;
    mx[u]=0;
    for(int i=0;i<G[u].size();i++)
    {
        int v=G[u][i].first;
        if(!vis[v]&&v!=fa)
        {
            dfssize(v,u);
            size[u]+=size[v];
            mx[u]=max(mx[u],size[v]);
        }
    }
}
void dfsroot(int r,int u,int fa)//r为子树的根 
{
    if(size[r]-size[u]>mx[u])//子树的其余节点 
        mx[u]=size[r]-size[u];
    if(mx[u]<MIN)//root为重心
        MIN=mx[u],root=u;
    for(int i=0;i<G[u].size();i++)
    {
        int v=G[u][i].first;
        if(!vis[v]&&v!=fa)
            dfsroot(r,v,u);
    }
}
void dfsdis(int u,int fa,int d)
{
    dis[num++]=d;
    for(int i=0;i<G[u].size();i++)
    {
        int v=G[u][i].first;
        int w=G[u][i].second;
        if(!vis[v]&&v!=fa)
            dfsdis(v,u,d+w);
    }
}
int cal(int r,int w)
{
    int ret=0;
    num=0;
    dfsdis(r,r,w);
    sort(dis,dis+num);
    int L=0,R=num-1;
    while(L<R)
    {
        while(dis[L]+dis[R]>k&&L<R)R--;
        ret+=R-L;    
        L++;
    }
    return ret;
}
void dfs(int u)
{
    MIN=n;
    dfssize(u,u);
    dfsroot(u,u,u);
    int Grivate=root;
    ans+=cal(Grivate,0);
    vis[root]=1;
    for(int i=0;i<G[Grivate].size();i++)
    {
        int v=G[Grivate][i].first;
        int w=G[Grivate][i].second;
        if(!vis[v])
        {
            ans-=cal(v,w);
            dfs(v);
        }
    }
}
int main()
{
    while(scanf("%d%d",&n,&k)!=EOF,n||k)
    {
        ans=0; 
        for(int i=1;i<=n;i++)
        {
            G[i].clear();
            vis[i]=0;
        }
        for(int i=1,u,v,w;i<n;i++)
        {
            scanf("%d%d%d",&u,&v,&w);
            G[u].push_back({v,w});
            G[v].push_back({u,w});
        }
        dfs(1);
        printf("%d
",ans);
    }
    return 0;
}