POI2014 HOT-Hotels

题目链接。

Description

给定一棵树，在树上选 3 个点，要求两两距离相等，求方案数。

Solution

因为 \(n \le 5000\)，可以 \(O(n ^ 2)\) 预处理两两之间的距离。

考虑是一颗树，而且这 \(3\) 个点两两距离相等，所以必然他们有一个汇聚点。且以这个汇聚点为根的话，三个点分的 \(lca\) 就是这个根。并且三个点到这个汇聚点的距离相等，而且两两的距离就是 \(点到汇聚点的距离 \times 2\)。这个东西容斥一下就好了，先枚举会聚点，然后把根删掉分成的每个联通快单独处理，每个联通快中，设距离为 \(d\) 的有 \(x\) 个，所有联通快中有 \(y\) 个。

• 先加上所有的 \(\dfrac{y \times (y - 1) \times (y - 2)}{6}\)
• 减掉两个在同一个子树的 \(x \times (x - 1) * (y - x)\)
• 减掉三个在同一个子树的 \(\dfrac{x \times (x - 1) \times (x - 2)}{6}\)

Tips

这题卡常，随便优化一下。

#include <iostream>
#include <cstdio>
#include <vector>
#define rint register int
using namespace std;

typedef long long LL;

const int N = 5005;

int n;

int head[N], cnt[N], s[N], q[N], d[N], rt, numE = 0;
vector<int> g[N];

bool vis[N];

LL ans = 0;

struct E{
	int next, v;
} e[N << 1];

void inline add(int u, int v) {
	e[++numE] = (E) { head[u], v };
	head[u] = numE;
}

void inline bfs(int x) {
	int hh = 0, tt = -1;
	q[++tt] = x; vis[x] = true;
	cnt[1]++;
	s[1]++; d[x] = 1;
	while (hh <= tt) {
		rint u = q[hh++];
		for (rint i = head[u]; i; i = e[i].next) {
			rint v = e[i].v;
			if (!vis[v]) {
				vis[v] = true;
				d[v] = d[u] + 1;
				cnt[d[v]]++;
				s[d[v]]++;
				q[++tt] = v;
			}
		}
	}
	for (rint i = 0; i <= tt; i++) {
		if (s[d[q[i]]]) {
			g[d[q[i]]].push_back(s[d[q[i]]]);
			s[d[q[i]]] = 0;
		}
	}
}
int main() {
	scanf("%d", &n);
	for (int i = 1, u, v; i < n; i++) 
		scanf("%d%d", &u, &v), add(u, v), add(v, u);
	for (rt = 1; rt <= n; rt++) {
		d[rt] = 0;
		for (rint j = 1; j <= n; j++) cnt[j] = 0, d[j] = 0, vis[j] = false, g[j].clear();
		vis[rt] = true;
		for (rint i = head[rt]; i; i = e[i].next) {
			bfs(e[i].v);
		}
		for (rint j = 1; j <= n; j++) {
			ans += (LL)cnt[j] * (cnt[j] - 1) * (cnt[j] - 2) / 6;
			for (rint i = 0; i < g[j].size(); i++) {
				rint x = g[j][i];
				ans -= (LL)x * (x - 1) / 2 * (cnt[j] - x) + (LL)x * (x - 1) * (x - 2) / 6;
			}
		}
	}

	printf("%lld\n", ans);
	return 0;
}