poj 1741 楼教主男人八题之一:树分治

http://poj.org/problem?id=1741

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
/**
poj 1741  楼教主男人八题之一:树分治
题目大意:给定一个带权树,问有多少对点之间的距离之和不超过k
解题思路:典型的树分治算法。(转自ACMonster)
           最容易想到的算法是:从每个点出发遍历整棵树,统计数对个数。
           由于时间复杂度O(N^2),明显是无法满足要求的。

           对于一棵有根树, 树中满足要求的一个数对所对应的一条路径,必然是以下两种情况之一:
           1、经过根节点
           2、不经过根节点,也就是说在根节点的一棵子树中
           对于情况2,可以递归求解,下面主要来考虑情况1。

           设点i的深度为Depth[i],父亲为Parent[i]。
           若i为根,则Belong[i]=-1,若Parent[i]为根,则Belong[i]=i,否则Belong[i]=Belong[Parent[i]]。
           这三个量都可以通过一次BFS求得。
           我们的目标是要统计:有多少对(i,j)满足i<j,Depth[i]+Depth[j]<=K且Belong[i]<>Belong[j]

           如果这样考虑问题会变得比较麻烦,我们可以考虑换一种角度:
           设X为满足i<j且Depth[i]+Depth[j]<=K的数对(i,j)的个数
           设Y为满足i<j,Depth[i]+Depth[j]<=K且Belong[i]=Belong[j]数对(i,j)的个数
           那么我们要统计的量便等于X-Y

           求X、Y的过程均可以转化为以下问题:
           已知A[1],A[2],...A[m],求满足i<j且A[i]+A[j]<=K的数对(i,j)的个数

           对于这个问题,我们先将A从小到大排序。
           设B[i]表示满足A[i]+A[p]<=K的最大的p(若不存在则为0)。我们的任务便转化为求出A所对应的B数组。那么,若B[i]>i,那么i对答案的贡献为B[i]-i。
           显然,随着i的增大,B[i]的值是不会增大的。利用这个性质,我们可以在线性的时间内求出B数组,从而得到答案。

           综上,设递归最大层数为L,因为每一层的时间复杂度均为“瓶颈”——排序的时间复杂度O(NlogN),所以总的时间复杂度为O(L*NlogN)

           然而,如果遇到极端情况——这棵树是一根链,那么随意分割势必会导致层数达到O(N)级别,对于N=10000的数据是无法承受的。因此,我们在每一棵子树中选择“最优”的点分割。所谓“最优”,是指删除这个点后最大的子树尽量小。这个点可以通过树形DP在O(N)时间内求出,不会增加时间复杂度。这样一来,即使是遇到一根链的情况时,L的值也仅仅是O(logN)的。

           因此,改进后算法时间复杂度为O(Nlog^2N),可以AC。
*/
#include <stdio.h>
#include <algorithm>
#include <iostream>
#include <string.h>
using namespace std;
const int maxn=10015;
const int INF=0x3f3f3f3f;
int head[maxn],ip;
int siz[maxn],can[maxn],lst[maxn],d[maxn],fa[maxn];
int ans,tl,n,k,l1,l2;

void init()
{
    memset(head,-1,sizeof(head));
    memset(can,0,sizeof(can));
    ip=0;
    ans=0;
}

struct note
{
    int v,w,next;
} edge[maxn*2];

void addedge(int u,int v,int w)
{
    edge[ip].v=v,edge[ip].w=w,edge[ip].next=head[u],head[u]=ip++;
}

void dfs1(int u,int pre)
{
    siz[u]=1;
    lst[++tl]=u;
    for(int i=head[u]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==pre||can[v])continue;
        dfs1(v,u);
        fa[v]=u;
        siz[u]+=siz[v];
    }
}
int getroot(int u,int pre)
{
    tl=0;
    dfs1(u,pre);
    int pos,tmp=INF,d,y;
    for(int i=1; i<=tl; i++)
    {
        d=0,y=lst[i];
        for(int p=head[y]; p!=-1; p=edge[p].next)
        {
            int v=edge[p].v;
            if(v==fa[y]||can[v])continue;
            d=max(d,siz[v]);
        }
        if(y!=u)
            d=max(d,siz[u]-siz[y]);
        if(d<tmp)pos=y,tmp=d;
    }
    return pos;
}

void dfs2(int u,int pre,int dis)
{
    lst[++l1]=u;
    d[u]=dis;
    for(int i=head[u]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==pre||can[v])continue;
        dfs2(v,u,dis+edge[i].w);
    }
}

int getans(int *a,int l,int r)
{
    int j=r,ret=0;
    for(int i=l; i<=r; i++)
    {
        while(d[a[i]]+d[a[j]]>k&&j>i)j--;
        ret+=j-i;
        if(j==i)break;
    }
    return ret;
}

inline bool cmp(int i,int j)
{
    return d[i]<d[j];
}
void work(int u,int pre)
{
    int root=getroot(u,pre);
    l1=l2=0;
    for(int i=head[root]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].v;
        if(can[v]==0)
        {
            l2=l1;
            dfs2(v,root,edge[i].w);
            sort(lst+l2+1,lst+l1+1,cmp);
            ans-=getans(lst,l2+1,l1);
        }
    }
    lst[++l1]=root,d[root]=0;
    sort(lst+1,lst+l1+1,cmp);
    ans+=getans(lst,1,l1);
    can[root]=1;
    for(int i=head[root]; i!=-1; i=edge[i].next)
    {
        int v=edge[i].v;
        if(can[v]==0)
            work(v,root);
    }
}
int main()
{
    while(~scanf("%d%d",&n,&k))
    {
        if(n==0&&k==0)break;
        init();
        for(int i=0; i<n-1; i++)
        {
            int x,y,z;
            scanf("%d%d%d",&x,&y,&z);
            addedge(x,y,z);
            addedge(y,x,z);
        }
        work(1,0);
        printf("%d\n",ans);
    }
    return 0;
}


你可能感兴趣的:(poj 1741 楼教主男人八题之一:树分治)