树的点分治 模板（POJ 1741)

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 22767		Accepted: 7527

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

题意：一棵有n个节点的树，每条边有个权值代表相邻2个点的距离，要求求出所有距离不超过k的点对(u,v)

题解：树分治
假设树以root为根节点，那么满足要求的点对有2种情况：
①路径经过root且dis(u,v)<=k
②路径不经过root，即其路径的最高点为子树上某一节点

对于第②种情况可以通过递归求解，这里只讨论第一种情况
该如何求解路径经过root且dis(u,v)<=k的合法点对数呢？

设dir[u]为u到根节点root的距离，那么只有满足dir[u]+dir[v]<=k且LCA(u,v)==root的点对才是合法的，
设cnt1=树中所有dis(u,root)+dis(root,v)<=k的点对数，cnt2=同属于root的某棵子树内的合法点对数
那么以root为根的树种合法点对数为：ans=cnt1-cnt2
找出有多少个dir[u]+dir[v]的方法很简单：只需要排序后扫一遍即可。

总结一下算法的过程：
①计算以u为根的树种每棵子树的大小
②根据子树大小找出树的重心root(以树的重心为根的树，可以使其根的子树中节点最多的子树的节点最少)
③以root为根，计算树中每个点到root的距离dir
④计算树中所有满足dir[u]+dir[v]<=k的点对数cnt1
⑤计算以root的子节点为根的子树中，满足dir[u]+dir[v]<=k的点对数cnt2
⑥ans+=cnt1-cnt2
注意：每次计算完cnt1后，要将vis[root]=1，这样就可以将一棵树分解成若干棵子树

#include <cstdio>  
#include <cstring>  
#include <iostream>  
#include <vector>  
#include <algorithm>  
using namespace std;  
const int MX = 1e4 + 5;  
  
struct Edge {  
    int v, w, nxt;  
} E[MX * 2];  
  
int n, k, root, Max, ans;  
vector <int> dis;  
int sz[MX], maxv[MX], head[MX], tot;  
bool vis[MX];  
  
void init() {  
    memset(vis, false, sizeof(vis));  
    memset(head, -1, sizeof(head));  
    tot = 0;  
}  
void add(int u, int v, int w) {  
    E[tot].v = v;  
    E[tot].w = w;  
    E[tot].nxt = head[u];  
    head[u] = tot++;  
}  
  
  
void dfs_size(int u, int fa) {  //求以U为根的子树中每个节点的子树大小
    sz[u] = 1; maxv[u] = 0;  
    for (int i = head[u]; ~i; i = E[i].nxt) {  
        int v = E[i].v;  
        if (vis[v] || v == fa) continue;  
        dfs_size(v, u);  
        sz[u] += sz[v];  
        maxv[u] = max(maxv[u], sz[v]);  
    }  
}  
  
void dfs_root(int r, int u, int pre) {  // 找出以u为根的子树的重心  
    maxv[u] = max(maxv[u], sz[r] - sz[u]);  
    if (Max > maxv[u]) {  //找到最大子树最小的节点即重心
        Max = maxv[u];  
        root = u;  
    }  
    for (int i = head[u]; ~i; i = E[i].nxt) {  
        int v = E[i].v;  
        if (v == pre || vis[v]) continue;  
        dfs_root(r, v, u);  
    }  
}  
  
void dfs_dis(int u, int fa, int dir) {  //算出每个节点到根节点的距离,并把它插入数组dis
    dis.push_back(dir);  
    for (int i = head[u]; ~i; i = E[i].nxt) {  
        int v = E[i].v, w = E[i].w;  
        if (vis[v] || v == fa) continue;  
        dfs_dis(v, u, dir + w);  
    }  
}  
  
int cal(int rt, int d) {  //算出以ri为根节点的子树中，dis(u,rt)+dis(rt,v)<=k的(u,v)的个数
    dis.clear();  
    dfs_dis(rt, -1, d);  
    sort(dis.begin(), dis.end());  //因为dis数组是无序的
    int i = 0, j = dis.size() - 1, ret = 0;  
    while (i < j) {  //扫一遍得出有多少对u,v，距离小于k，（对于每一个点i求满足条件的点有多少）
        while (dis[i] + dis[j] > k && i < j) j--;  
        ret += j - i;  
        i++;  
    }  
    return ret;  
}  
  
void DFS(int u) {  //计算以u为根节点的子树中，满足条件的所有(u,v)对的个数
    Max = n;  
    dfs_size(u, -1);  //这两步都是为了算出以u为根节点的树的重心
    dfs_root(u, u, -1);  
    int rt = root;  
    ans += cal(rt, 0);  //以u作为根的树所有满足dir[u]+dir[v]<=k且经过根节点的点对数cnt1
    vis[rt] = 1;  
    for (int i = head[rt]; ~i; i = E[i].nxt) {  //求以root的子节点为根的子树中，满足dir[u]+dir[v]<=k的点对数cnt2
        int v = E[i].v, w = E[i].w;  
        if (vis[v]) continue;  
		//下面这里是算只有以v为子树的根（即节点都是u的左子树或右子树的点时），满足条件的情况的点，即多余不需要的情况，w表示这些都是到根节点u的距离，为了与cal(rt, 0)保持一致，找到多余的情况
		//即dis(u,rt)+dis(rt,v)<=k,且u,v∈u的左(右)子树
        ans -= cal(v, w);    //cnt2    这里只需要计算经过子根节点v的dis(u,v)<=k但实际不需要经过根节点u的个数，因为这些点将在下面的DFS中计算，这里把他删掉，以免重复。
		//1.路径经过root且dis(u,v)<=k
		//因为任何(u,v)一定经过根节点，这样只要算每一个节点作为根节点，算出经过他的(u,v)的个数。如果不按这样算的话，用DFS算的话就会重复
        DFS(v);  //算出以v为根节点的(u,v)的个数  2.路径不经过root，即其路径的最高点为子树上某一节点
    }  
}  
  
int main() {  
    //freopen("in.txt","r",stdin);  
    while (scanf("%d%d", &n, &k), n || k) {  
        init();  
        for (int i = 1; i < n; i++) {  
            int u, v, w;  
            scanf("%d%d%d", &u, &v, &w);  
            add(u, v, w); add(v, u, w);  
        }  
        ans = 0;  
        DFS(1);  
        printf("%d\n", ans);  
    }  
    return 0;  
}