poj 1741 Tree （树的分治）

Tree

Time Limit: 1000MS Memory Limit: 30000K
Total Submissions: 20776 Accepted: 6803

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

这是一道树的分治题目，关于树的分治可以看漆子超《分治算法在树的路径问题中的应用》，里面详细讲到了这道题目的思路。
一条合法的路径，要么是经过根节点，要么是在子树中。我们只考虑经过根节点的情况，剩下的子树可以进行递归处理。我们首先计算出每个点到根节点的距离，如果depth[i]+depth[j]<=k&&i j不在同一颗子树中，我们就将它加入结果的统计。所以结果即为depth[i]+depth[j]<=k - depth[i]+depth[j]<=k&&i j在同一颗子树中。显然，ij在同一棵子树也是满足情况的，我们之所以要减去，是因为在接下来会对每棵子树进行相同的处理，这条路径一定会经过某个子树的根节点，如果不减去就会造成重复计算。
分治算法中要首先确定一棵树的重心，先任选一个节点作为根，把无根树变为有根树，然后进行dfs求出s[i]即以这个点为根节点子树的大小，我们在用mx[i]记录在i的所有子树中size最大的子树的size，选取mx[i]最小的点即为树的重心。（树的重心即为最大子树最小的点）

而cala函数即为计算相应的点对数，首先求出子树中所有点到这个点的距离（getdepth函数），然后经过一个排序快速求出所有满足条件的点对，然后对于其所有子树中，求出所有满足条件的点对并相减就是以这个点为根时的答案。

之前一直TLE，后来经过discuss中的大佬的提醒才找到原因。
在不断求重心的过程中，msize和minsize的值都是要不断更新为s[i]，而不是始终为n，这样会导致求出来的重心错误，将复杂度由o(logn)退化为o(n)，导致超时。

以下是AC代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>
#include <algorithm>
#define maxn 10010
using namespace std;
struct node{
    int to,len;
    node(){};
    node(int _to,int _len):to(_to),len(_len) {};
};

vector<node> g[maxn];
vector<int> dep;
int n,k;
int root,msize,minsize;
int s[maxn],mx[maxn],vis[maxn],d[maxn];        //s用于存储每个节点及其所有子树的数目;mx数组用于存储将这个节点去掉后的子树中节点个数最大的一个，所以x[i]的值越小，此时的子树划分越均衡，这便是此时的重心。
int ans,num;
void getroot(int now,int fa)
{
    s[now]=1;
    mx[now]=0;
    for(int i=0;i<g[now].size();i++)
    {
        int v=g[now][i].to;
        if(v==fa||vis[v]) continue;
        getroot(v,now);
        s[now]+=s[v];
        mx[now]=max(mx[now],s[v]);     //在其所有子树中找出一个节点数最多的点
    }
    mx[now]=max(mx[now],msize-s[now]);      //另外一种情况就是总节点数减去这个点的size值，便是另外一种子树的情况。
    if(mx[now]<minsize) {
        minsize=mx[now];
        root=now;
    }
}

void getdep(int now,int fa,int len)
{
    d[num++]=len;
    for(int i=0;i<g[now].size();i++)
    {
        int v=g[now][i].to;
        if(v==fa||vis[v]) continue;
        getdep(v,now,len+g[now][i].len);
    }
}


int calc(int now,int len)
{
    num=0;
    int ret=0;
    getdep(now,0,len);
    sort(d,d+num);
    int i=0,j=num-1;
    while(i<j)         //很巧妙的方法……求出这些距离中相加满足条件的共有多少对。
    {
        while(d[i]+d[j]>k&&i<j) j--;
        ret+=j-i;
        i++;
    }
    return ret;
}

void solve(int now)
{
    ans+=calc(now,0);
    vis[now]=1;
    for(int i=0;i<g[now].size();i++)
    {
        int v=g[now][i].to;
        if(vis[v]) continue ;
        ans-=calc(v,g[now][i].len);//这里传进这个长度的参数，是为了算出所有到其父节点的距离，即和循坏之上的那个cala算出来的距离是相同的，便是为了求出只在这颗子树中，所有满足条件的点对。
        msize=minsize=s[v];
        getroot(v,root=0);
        solve(root);
    }
}

int main()
{
    while(scanf("%d%d",&n,&k)!=EOF)
    {
        if(n==0&&k==0) break;
        for(int i=0;i<=n;i++)
            g[i].clear();
        for(int i=0;i<n-1;i++)
        {
            int u,v,l;
        //    cin>>u>>v>>l;
            scanf("%d%d%d",&u,&v,&l);
            g[u].push_back(node(v,l));
            g[v].push_back(node(u,l));
        }
        ans=0;
        msize=n;
        minsize=n;
        memset(vis,0,sizeof(vis));
        getroot(1,root=0);
        solve(root);
        cout<<ans<<endl;
    }
    return 0;
}