poj1741树分治

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 11912		Accepted: 3755

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

Source

LouTiancheng@POJ

题意：给定一棵树和距离k，求距离小于k节点对的个数

解法：在树上做分治。按照重心分解，分成在同一颗子树上和不同子树上两种情况讨论，不过要注意去重

代码：

//
//  main.cpp
//  poj1741
//
//  Created by 蘇與軒 on 15/2/9.
//  Copyright (c) 2015年 蘇與軒. All rights reserved.
//
#include <iostream>
#include <cstdio>
#include <cstring>
#include <queue>
#include <vector>
#include <cmath>
#include <cstdlib>
#include <string>
#include <map>
#include <set>
#include <algorithm>
#include <functional>
#define rep(i,a,b) for (int i=a;i<((b)+1);i++)
#define Rep(i,a,b) for (int i=a;i>=b;i--)
#define foreach(e,x) for (__typeof(x.begin()) e=x.begin();e!=x.end();e++)
#define mid ((l+r)>>1)
#define lson (k<<1)
#define rson (k<<1|1)
#define MEM(a,x) memset(a,x,sizeof a)
using namespace std;
const int N=2000050;
typedef pair<int, int> pii;
typedef long long ll;
int pnt[N],head[N],nxt[N],dist[N],n,k,cnt,ans,size,f[N],tsize[N],root,d[N];
bool vis[N];
vector<int> dis;
void addedge(int u,int v,int c) {
    pnt[cnt]=v;nxt[cnt]=head[u];dist[cnt]=c;head[u]=cnt++;
    pnt[cnt]=u;nxt[cnt]=head[v];dist[cnt]=c;head[v]=cnt++;
}
void subtree(int u,int p) {
    int v;
    tsize[u]=1;f[u]=0;
    for (int i=head[u];i!=-1;i=nxt[i]) {
        v=pnt[i];
        if (v==p||vis[v]) continue;
        subtree(v, u);
        tsize[u]+=tsize[v];
        f[u]=max(f[u],tsize[v]);
    }
    f[u]=max(f[u],size-f[u]);
    if (f[u]<f[root])   root=u;
}
void getpath(int u,int p) {
    dis.push_back(d[u]);
    int v;
    tsize[u]=1;
    for (int i=head[u];i!=-1;i=nxt[i]) {
        v=pnt[i];
        if (v!=p&&!vis[v]) {
            d[v]=d[u]+dist[i];
            getpath(v, u);
            tsize[u]+=tsize[v];
        }
    }
}
int cal(int u,int init) {
    dis.clear();d[u]=init;
    getpath(u,0);
    sort(dis.begin(),dis.end());
    int ret=0;
    for (int l=0,r=(int)dis.size()-1;l<r;)
        if (dis[l]+dis[r]<=k)   ret+=r-l++;
        else r--;
    return ret;
}
void solve(int u) {
    vis[u]=true;
    int v;
    for (int i=head[u];i!=-1;i=nxt[i]) {
        v=pnt[i];
        if (!vis[v]) {
            f[0]=size=tsize[u];
            ans-=cal(v,dist[i]);
            subtree(v, root=0);
            solve(root);
        }
    }
    ans+=cal(u,0);
    vis[u]=false;
}
int main(int argc, const char * argv[]) {
    while (scanf("%d%d",&n,&k)!=EOF){
        if ((n|k)==0) break;
        MEM(head,-1);MEM(nxt,-1);
        rep(i,1,n-1){
            int u,v,c;
            scanf("%d%d%d",&u,&v,&c);
            addedge(u,v,c);
        }
        size=f[0]=n;
        subtree(1,root=0);
        ans=0;
        solve(root);
        printf("%d\n",ans);
    }
    return 0;
}