POJ1741——Tree（树上分治，树的重心）

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 17804		Accepted: 5818

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的点对的个数。

白书上的例题，说实话没有太明白。

感觉白书的代码比较有条理但还是过于冗长了。

#include <stdio.h>
#include <string.h>
#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <set>
#include <map>
#include <string>
#include <math.h>
#include <stdlib.h>
using namespace std;

#define MAXN 20010
struct edge
{
    int to,length;
};

int n,k;
vector <edge> G[MAXN];

bool centroid[MAXN];
int subtree_size[MAXN];


void init(){
    for(int i=0;i<=n;i++)
        G[i].clear();
    memset(centroid,0,sizeof(centroid));
    memset(subtree_size,0,sizeof(subtree_size));
}

int ans;

int com_sub(int v,int p){
    int c=1;
    for(int i=0;i<G[v].size();i++){
        int w=G[v][i].to;
        if(w==p||centroid[w])
            continue;
        c+=com_sub(G[v][i].to,v);
    }
    subtree_size[v]=c;
    return c;
}

pair <int,int> ser_cen(int v,int p,int t){
    pair <int,int> res=make_pair(INT_MAX,-1);
    int s=1,m=0;
    for(int i=0;i<G[v].size();i++){
        int w=G[v][i].to;
        if(w==p||centroid[w])
            continue;
        res=min(res,ser_cen(w,v,t));
        m=max(m,subtree_size[w]);
        s+=subtree_size[w];
    }
    m=max(m,t-s);
    res=min(res,make_pair(m,v));
    return res;
}

void en_path(int v,int p,int d,vector <int> &ds){
    ds.push_back(d);
    for(int i=0;i<G[v].size();i++){
        int w=G[v][i].to;
        if(w==p||centroid[w])
            continue;
        en_path(w,v,d+G[v][i].length,ds);
    }
}

int count_pair(vector <int> &ds){
    int res=0;
    sort(ds.begin(),ds.end());
    int j=ds.size();
    for(int i=0;i<ds.size();i++){
        while(j>0&&ds[i]+ds[j-1]>k)
            j--;
        res+=j-(j>i?1:0);
    }
    return res/2;
}

void _solve(int v){
    com_sub(v,-1);
    int s=ser_cen(v,-1,subtree_size[v]).second;
    centroid[s]=true;

    for(int i=0;i<G[s].size();i++){
        if(centroid[G[s][i].to])
            continue;
        _solve(G[s][i].to);
    }
    vector <int> ds;
    ds.push_back(0);
    for(int i=0;i<G[s].size();i++){
        if(centroid[G[s][i].to])
            continue;
        vector <int> tds;
        en_path(G[s][i].to,s,G[s][i].length,tds);
        ans-=count_pair(tds);
        ds.insert(ds.end(),tds.begin(),tds.end());
    }
    ans+=count_pair(ds);
    centroid[s]=false;
}

void solve(){
    ans=0;
    _solve(0);
    printf("%d\n",ans);
}


int main(){
    while(scanf("%d%d",&n,&k)){
        if(n==0&&k==0)
            break;
        init();
        for(int i=0;i<n-1;i++){
            int u,v,w;
            edge a;
            edge b;
            scanf("%d%d%d",&u,&v,&w);
            u--;
            v--;
            a.to=v;
            b.to=u;
            a.length=w;
            b.length=w;
            G[u].push_back(a);
            G[v].push_back(b);
        }
//        printf("%d\n",G[1].size());
//        printf("%d\n",G[2].size());
        solve();
    }
}