POJ 1741 Tree(树分治+容斥原理+树的重心)

时间:2020-12-20 04:24:43
Tree
Time Limit: 1000MS   Memory Limit: 30000K
Total Submissions: 25261   Accepted: 8413

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



        一些很神奇的算法杂合在一起的题目。
        大致题意:给你一个带边权的树,问你总共有多少对点距离在K以内,这里定义点对是无序的,即(a,b)和(b,a)算是同一对。
        其实呢,之前我在做HDU多校赛的一些题目的时候遇到过类似的树上统计的问题,其实这类做法有一个通用的名称——树分治。就是统计树上的东西,我们可以根据树的结构分治,对于一个节点,我可以把它的每个儿子分开治,把儿子当作单独的一个问题处理,然后最后合起来,再去除一些重复的东西即可。这题也正是这么做的。
        首先,我们不考虑重复的问题。对于以一个点为根的子树,我们可以线性计算出该点到达所有点的距离。知道距离之后,我们就可以用一种很经典又很神奇的方法线性求出所有距离小于等于K的点对数量。我们把所有的距离放到一个数组dist,然后从大到小排序,分别定义首尾指针l、r。然后如果dist[l]+dist[r]<=K,那么显然,以l为第一个点,以r或者r前面任意点为第二个点,这样的点对的距离肯定小于等于K,然后指针l往右移动;如果大于,则说明r大了,指针r往左移动。这样子的话相当于找了所有路径经过根的点对和一些最短路径不经过根但是绕远之后距离还是小于等于K的点对。而那多的部分就是我们接下来要处理。
        其次,再考虑这个多计算的部分。我们设我们刚刚计算出来的东西为cal(x),但是我们实际想要计算的只是最短路径经过根的点对。减去这些,我们只需要减去x所有的儿子对应的cal(y)即可,因为cal(y)同样包括了所有最短路径经过y的点对和一些不经过y的点对。对所有的儿子都这么处理,剩下的就只能是最短路径经过根x的点对了。现在想来,这个去重的过程不正是容斥原理吗。
        这样子在具体代码中,如果不算排序,相当于在dfs中,加了一个同等复杂度的另一个dfs,对总的复杂度只是常数的改变。然后加上了排序之后无非是多了个logN,理想的总复杂度是O(NlogNlogN)。但是实际不一定是那么理想,如果是链状的树,那么会退化到O(N^2)更准确点应该是O(N^2logN)。为了避免这个情况,又要用到一个比较经典同时又很巧妙的方法。由于我用的是分治,而且仅仅是计算距离,子树内部的变换不影响外部,所以在分治一个小的子树的时候,我可以任意选择起始的根。而这个根为了时间的高效,就选用了树的重心。这个同样可以线性求出。
        就这样,这题杂合了树分治、容斥原理和树的重心,可以说非常有意思。具体见代码:
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
#define N 10010

using namespace std;

int d[N],dist[N],son[N],f[N],n,k,ans,num,root;
struct Edge{int y,w;};
vector<Edge> g[N];
bool v[N];

void getroot(int x,int fa)						//找树的重心
{
    son[x]=1; f[x]=0;
    for(int i=0;i<g[x].size();i++)
    {
        int y=g[x][i].y;
        if (y==fa||v[y]) continue;
        getroot(y,x);
        son[x]+=son[y];
        f[x]=max(f[x],son[y]);
    }
    f[x]=max(num-son[x],f[x]);
    if (f[x]<f[root]) root=x;
}

void cal_dist(int x,int fa)						//求根到所有点的距离
{
    dist[++dist[0]]=d[x];
    for(int i=0;i<g[x].size();i++)
    {
        int y=g[x][i].y;
        int w=g[x][i].w;
        if (y==fa||v[y]) continue;
        d[y]=d[x]+w; cal_dist(y,x);
    }
}

int cal(int x,int w)
{
    dist[0]=0;int res=0;
    d[x]=w; cal_dist(x,root);
    int l=1,r=dist[0];
    sort(dist+1,dist+1+r);
    while(l<r)									//神奇的方法计算点对
        if (dist[l]+dist[r]<=k) res+=r-l,l++;
                                    else r--;
    return res;
}

void dfs(int x)
{
    v[x]=1;
    ans+=cal(x,0);
    for(int i=0;i<g[x].size();i++)
    {
        int y=g[x][i].y;
        int w=g[x][i].w;
        if (v[y]) continue;
        ans-=cal(y,w);								//容斥原理
        num=son[y]; root=0;
        getroot(y,0); dfs(root);					//分治子树的时候,选取子树的重心作为根,防止退化
    }
}

int main()
{
    while(~scanf("%d%d",&n,&k))
    {
        if (n+k==0) break;
        ans=root=0; num=n;
        memset(v,0,sizeof(v));
        memset(g,0,sizeof(g));
        for(int i=1;i<n;i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            g[u].push_back(Edge{v,w});
            g[v].push_back(Edge{u,w});
        }
        f[0]=N;
        getroot(1,0);
        dfs(root);
        printf("%d\n",ans);
    }
    return 0;
}