BZOJ1468 Tree <点分治>

Problem

Tree

$\mathrm{Time\;Limit:\;10\;Sec}$
$\mathrm{Memory\;Limit:\;64\;MB}$

Description

给你一棵树以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于$K$

Input

$N$（$N\le 40000$） 接下来$n-1$行边描述管道，按照题目中写的输入 接下来是$K$

Output

一行，有多少对点之间的距离小于等于$K$

Sample Input

7
1 6 13 
6 3 9 
3 5 7 
4 1 3 
2 4 20 
4 7 2 
10

Sample Output

5

标签：点分治

Solution

点分治板题。
对于每个点，我们计算穿过它的路径的条数。
这个值等于它的子树中所有深度之和小于$k$的点对数，可以双指针搞一搞。
但这样两个在同一儿子内的点对会算重，所以需要减去它的每个儿子的子树中深度之和小于$k-c$的点对数，其中$c$为它到此儿子的距离。这也可以双指针搞一搞。
对于每次分治，我们都要找到当前子树的重心，以保证复杂度为$O(n\log(n))$。

Code

#include <bits/stdc++.h>
#define MAX_N 40000
using namespace std;
int n, k, rt, s[MAX_N+5], h[MAX_N+5], d[MAX_N+5], tot, ans;
vector <int> G[MAX_N+5], E[MAX_N+5], dep;    bool mrk[MAX_N+5];
void addedge(int u, int v, int c) {G[u].push_back(v), E[u].push_back(c);}
void getrt(int u, int fa) {
    s[u] = 1, h[u] = 0;
    for (int i = 0; i < (int)G[u].size(); i++) {
        int v = G[u][i];	if (v == fa || mrk[v])	continue;
        getrt(v, u), s[u] += s[v], h[u] = max(h[u], s[v]);
    }
    h[u] = max(h[u], tot-s[u]);	if (h[u] < h[rt]) rt = u;
}
void getdep(int u, int fa) {
    dep.push_back(d[u]), s[u] = 1;
    for (int i = 0; i < (int)G[u].size(); i++) {
        int v = G[u][i], c = E[u][i];
        if (v == fa || mrk[v])	continue;
        d[v] = d[u]+c, getdep(v, u), s[u] += s[v];
    }
}
int calc(int u, int init) {
    dep.clear(), d[u] = init, getdep(u, 0), sort(dep.begin(), dep.end());	int ret = 0;
    for (int l = 0, r = (int)dep.size()-1; l < r; ) dep[l]+dep[r] <= k ? ret += r-l++ : r--;
    return ret;
}
void DFS(int u) {
    ans += calc(u, 0), mrk[u] = true;
    for (int i = 0; i < (int)G[u].size(); i++) {
        int v = G[u][i], c = E[u][i];	if (mrk[v])	continue;
        ans -= calc(v, c), h[0] = tot = s[v], getrt(v, rt = 0), DFS(rt);
    }
}
int main() {
    scanf("%d", &n);	for (int i = 1, u, v, c; i < n; i++) scanf("%d%d%d", &u, &v, &c), addedge(u, v, c), addedge(v, u, c);
    scanf("%d", &k);	h[0] = tot = n, getrt(1, 0);	DFS(rt);	printf("%d", ans);	return 0;
}