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_JUDGEfreopen("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;}