jzoj 1166. 树中点对距离

Description

给出一棵带边权的树，问有多少对点的距离<=\(len\)
100% \(2<=n<=10000,len<=maxlongint\)

Solution

这一题可以说是点分治的模板题了。
我们按照套路，先求重心，在计算答案。
如何计算答案？
设当前点为\(x\)
我们先\(O(n)\)搜一遍求出当前树的每个点的深度。
分类讨论：
经过\(x\)，那么只需满足\(dep[a] + dep[b] <= len\)即可。
不经过\(x\)，那么这个可以在那一颗子树中计算到。
所以，我们需要容斥，将不经过\(x\)的减掉。
由于他们到\(x\)还有相同的一段路径（设长\(length\)），所以要将他们的\(dep\)之和加上\(2*length\)。

Code

#include <cstdio>
#include <algorithm>
#define N 10010
using namespace std;
struct node{int v, fr, l;}e[N << 1];
int n, len, rt, cnt = 0, ans = 0, tot;
int tail[N], siz[N], son[N], dep[N], size = 0;
bool bz[N];

inline int read()
{
	int x = 0; char c = getchar();
	while (c < '0' || c > '9') c = getchar();
	while (c >= '0' && c <= '9') x = (x << 1) + (x << 3) + (c ^ 48), c = getchar();
	return x;
}

void add(int u, int v, int l) {e[++cnt] = (node){v, tail[u], l}; tail[u] = cnt;}

void getrt(int x, int fa)
{
	siz[x] = 1; son[x] = 0;
	for (int p = tail[x], v; p; p = e[p].fr)
	{
		v = e[p].v;
		if (v == fa || bz[v]) continue;
		getrt(v, x);
		if (son[x] < siz[v]) son[x] = siz[v];
		siz[x] += siz[v];
	}
	if (son[x] < tot - siz[x]) son[x] = tot - siz[x];
	if (son[x] < son[rt]) rt = x;
}

void getdeep(int x, int fa, int deep)
{
	dep[++size] = deep;
	for (int p = tail[x], v; p; p = e[p].fr)
	{
		v = e[p].v;
		if (v == fa || bz[v]) continue;
		getdeep(v, x, deep + e[p].l);
	}
}

int cal(int x, int num)//calculation
{
	size = 0, getdeep(x, 0, num);
	sort(dep + 1, dep + size + 1);
	int l = 1, r = size, s = 0;
	while (l < r)
	{
		while (dep[l] + dep[r] <= len && l < r) s += r - l, l++;
		while (dep[l] + dep[r] > len && l < r) r--;
	}
	return s;
}

void solve(int x)
{
//	printf("%d\n", x);
	bz[x] = 1;
	ans += cal(x, 0); 
	for (int p = tail[x], v; p; p = e[p].fr)
	{
		v = e[p].v;
		if (bz[v]) continue;
		ans -= cal(v, e[p].l);
		tot = siz[v], rt = 0, getrt(v, 0);
		solve(rt);
	}
}

int main()
{
	freopen("distance.in", "r", stdin);
	freopen("distance.out", "w", stdout);
	n = read(), len = read(); son[0] = 1e9;
	for (int i = 1, u, v, l; i < n; i++)
		u = read(), v = read(), l = read(), add(u, v, l), add(v, u, l);
	tot = n, rt = 0, getrt(1, 0);
	solve(rt);
	printf("%d\n", ans); 
	return 0;
}