POJ 1741 Tree (树分治模板题)

点击打开链接

题意：给出一棵树，求出两点之间距离小于等于k的点对数目。

解法：树的考查十分常见，原因就在于树有一般图不具备的约束条件。

如果这是一个一般图，我们想的是，做一次LCA，然后暴力两个点，看是否符合，复杂度达到O（n²），对于题目有1e4个点而且还是多组输入显然是无法接受的。

这时候我们就要想到树分治的思想。

求两个点的距离是否小于等于k，可以分为两种情况。

我们假定一个根节点为root，那么两个点距离小于等于k，这条路有可能经过root，也有可能不经过root。

对于上面两种情况。

1.如果经过root，那就简单了，我们直接求出root到每个点的距离，排序，然后O(n)扫一遍就可以知道有多少个点是符合答案的了。

比如样例中，假设root = 1，那么就可以得到1 2 2 3这些距离，令l = 0， r = 3，O（n）扫一遍，得到的答案是1 2,1 2，1 3,2 2。所以有四个点对是符合条件的。但是事实上，节点3到节点5是没有经过root的，所以我们需要把这种情况给去掉（从root枚举连接的点v，走他的子树，初始化v的距离为 root到v的距离，找出仍然符合条件的数目，具体看代码），所以经过root=1的点对符合条件的是三种。

2.如果不经过root，其实每条路肯定会经过某个点的，所以直接从root递归下去指定新的root即可。

如果root选定得不好，很容易使得复杂度退化为n²，所以树的分治是有求重心的部分在里面的，令重心为root，可以保证树的深度尽量小。

时间复杂度O（nlog²n）。复杂度证明可以搜《分治算法在树的路径问题中的应用》。

代码如下：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<cmath>
#include<ctime>
#include<algorithm>
using namespace std;
const int maxn = 1e4 + 5;
const int INF = 0x3f3f3f3f;
int tot, head[maxn], to[maxn << 1], cost[maxn << 1], nx[maxn << 1];
int n, k, root, sum, ans, num;//sum存放子树点数 
int f[maxn], vis[maxn], son[maxn], deep[maxn];
//f存点i的最大的子树节点数目 
void add_edge(int u,int v,int val) {
    to[tot] = v, cost[tot] = val, nx[tot] = head[u], head[u] = tot++;
    swap(u, v);
    to[tot] = v, cost[tot] = val, nx[tot] = head[u], head[u] = tot++;
}
//找树的重心 
void getroot(int u, int fa) {
    son[u] = 1;
	f[u] = 0;
    for(int i = head[u]; ~i; i = nx[i]) {
    	int v = to[i];
        if(v == fa || vis[v])  
			continue;
        getroot(v, u);
        son[u] += son[v]; //记录子树有多少个节点 
        f[u] = max(f[u], son[v]);
    }
    f[u] = max(f[u], sum - son[u]);
    if(f[u] < f[root])
		root = u;
}
//计算每个点到树的重心的距离存于deep中
//deep[0]表示有多少个元素 
void getdeep(int u, int fa, int last_dis) {
    for(int i = head[u]; ~i; i = nx[i]) {
    	int v = to[i];
        if(v == fa || vis[v])
			continue;
        deep[num++] = last_dis + cost[i];
        getdeep(v, u, deep[num - 1]);
    }
}
//计算有多少个点对长度小于等于k 
int cal(int u, int val) {
	num = 0;
	deep[num++] = val;
    getdeep(u, 0, deep[num - 1]);
    sort(deep, deep + num);
    int l = 0, r = num - 1, tmp = 0;
    while(l < r) { //扫一遍即可 
        if(deep[l] + deep[r] <= k) { //如果成立，那么 
			tmp += r - l;//以l为首到r的两个点之间都会成立 
			l++;
		} else //否则收缩最大的范围 
			r--;
    }
    return tmp;
}

void solve(int u) {
    ans += cal(u, 0);//计算答案
    vis[u] = 1;
    for(int i = head[u]; ~i; i = nx[i]) {
    	int v = to[i];
        if(vis[v])
			continue;
        ans -= cal(v, cost[i]);//计算不符合题意的答案
        sum = son[v];
        root = 0;
        getroot(v, 0); //找出子树的重心 
        solve(root); //递归解决 
    }
}

int main() {
#ifndef ONLINE_JUDGE
	freopen("in.txt", "r", stdin);
#endif
    while(scanf("%d%d", &n, &k) != EOF) {
    	if(n == 0 && k == 0)
			break;
        ans = 0, root = 0, tot = 0;
        memset(vis, 0, sizeof(vis));
        memset(head, -1, sizeof(head));
        for(int i = 1, u, v, val; i < n; i++) {
            scanf("%d%d%d", &u, &v, &val);
            add_edge(u, v, val);
        }
        f[0] = INF;//root=0，用来比较找出子树节点最小的重心 
		sum = n;
        getroot(1, 0);
        solve(root);
        printf("%d\n", ans);
    }
    return 0;
}