POJ1741树分治-点分治

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

#include
#include
#include
#include
using namespace std;
int n,k;
struct Edge
{
    int fro,to,val,next;
}ed[10005*2];
int cnt,head[10005];
int vis[10005];//这个点有没有被砍掉
int siz[10005];//以i为根节点的子树的大小
int maxn[10005];//以i为根节点的所有子树中最大的子树大小 
vector<int>dis;
int ans,root,Max;// root表示最后选择的虚根，Max表示以root为根的子树中最大的子树 
int size;
void add(int u,int v,int w)
{
    ed[cnt].fro=u;
    ed[cnt].to=v;
    ed[cnt].val=w;
    ed[cnt].next=head[u];
    head[u]=cnt++;
}
void get_size(int u,int fa)//求出size数组和maxn数组 
{
    siz[u]=1; 
    maxn[u]=0;
    for(int i=head[u];~i;i=ed[i].next)
    {
        int v=ed[i].to;
        if(v!=fa&&!vis[v])
        {
            get_size(v,u);
            siz[u]+=siz[v];
            maxn[u]=max(maxn[u],siz[v]);
        }
    }


}
void get_root(int u,int fa)//找出以u为根节点的子树的重心 maxn[]最小的结点是重心 
{
    maxn[u]=max(maxn[u],size-siz[u]);//size表示最初的u为根节点的子树的大小，不在递归中改变 
    if(Max>maxn[u])
    {
        Max=maxn[u];
        root=u;
    }
    for(int i=head[u];~i;i=ed[i].next)
    {
        int v=ed[i].to;
        if(v!=fa&&!vis[v])
        {
            get_root(v,u);
        }
    }
}

//如果满足i->fa->j<=k,肯定有i->u->j<=k 
//存在 i->fa->j>k，使得 i->u->j<=k ,这种情况下的ij不会被重复计算
//求重复计算的i j 要满足 i->u->fa+j->u->fa<=k 
void get_dis(int u,int fa,int d)//找出以u为根节点的子树中每个点到fa的距离 
{
    dis.push_back(d);
    for(int i=head[u];~i;i=ed[i].next)
    {
        int v=ed[i].to;
        int w=ed[i].val;
        if(v!=fa&&!vis[v])
        {
            get_dis(v,u,d+w);
        }
    } 
}

int cal(int u,int d)//返回以u为根节点的子树中 i->u->j的距离不超过k的(i,j)对数 
{
    dis.clear();
    get_dis(u,-1,d);//d?
    sort(dis.begin(),dis.end());
    int i=0,j=dis.size()-1,sum=0;
    while(iwhile(dis[i]+dis[j]>k&&ireturn sum;
}
void dfs(int u)
{
    Max=n;
    get_size(u,-1);
    size=siz[u];
    get_root(u,-1);//找到以u为根的树中重心为root，重新以root为根 
    ans+=cal(root,0);//以root为根的子树中 i->root->j<=k的ij有多少 
    vis[root]=1;
    for(int i=head[root];~i;i=ed[i].next)
    {
        int v=ed[i].to;
        int w=ed[i].val;
        if(!vis[v])
        {
            ans-=cal(v,w);//减去 i->v->j（即不经过root的i->j）<=k 的ij 
            //size=siz[v];
            dfs(v); 
        }
    }

} 
int main()
{
    while(~scanf("%d%d",&n,&k)&&(n||k))
    {
        for(int i=1;i<=n;i++)
        head[i]=-1,vis[i]=0;
        cnt=0;
        Max=n;
        root=1;
        ans=0;
        for(int i=1;iint u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);add(v,u,w);
        }
        dfs(1);
        printf("%d\n",ans);
    }
    return 0;
}