P3565 [POI2014]HOT-Hotels

\(\color{#0066ff}{ 题目描述 }\)

有一个树形结构，每条边的长度相同，任意两个节点可以相互到达。选3个点。两两距离相等。有多少种方案？

\(\color{#0066ff}{输入格式}\)

第一行：n

然后是n-1条边

\(\color{#0066ff}{输出格式}\)

方案数

\(\color{#0066ff}{输入样例}\)

7
1 2
5 7
2 5
2 3
5 6
4 5

\(\color{#0066ff}{输出样例}\)

5

\(\color{#0066ff}{数据范围与提示}\)

\(1\leq n \leq 5000\)

\(\color{#0066ff}{ 题解 }\)

不难发现，合法的三个点一定是一个电风扇形qwq

放在树上，如果当前点是根，那么不同子树中的某一层，每个子树中选出一个点（当前点为LCA）一定是成立的，而且只有这种情况

所以枚举每个点为根，dfs每棵子树，用桶统计每个点的深度

用两个数组维护已经搜过的子树选一个点， 选两个点的方案，更新就行啦

#include<bits/stdc++.h>
#define LL long long
LL in() {
	char ch; LL x = 0, f = 1;
	while(!isdigit(ch = getchar()))(ch == '-') && (f = -f);
	for(x = ch ^ 48; isdigit(ch = getchar()); x = (x << 1) + (x << 3) + (ch ^ 48));
	return x * f;
}
struct node {
	int to;
	node *nxt;
	node(int to = 0, node *nxt = NULL):to(to), nxt(nxt) {}
	void *operator new (size_t) {
		static node *S = NULL, *T = NULL;
		return (S == T) && (T = (S = new node[1024]) + 1024), S++;
	}
};
const int maxn = 5050;
int c[maxn], c2[maxn], t[maxn];
int du[maxn];
int n, maxdep;
LL ans;
node *head[maxn];
void add(int from, int to) {
	head[from] = new node(to, head[from]);
}
void dfs(int x, int fa, int dep) {
	t[dep]++;
	maxdep = std::max(maxdep, dep);
	for(node *i = head[x]; i; i = i->nxt)
		if(i->to != fa)
			dfs(i->to, x, dep + 1);
}
int main() {
	n = in();
	int x, y;
	for(int i = 1; i < n; i++) {
		x = in(), y = in();
		add(x, y), add(y, x);
		du[x]++, du[y]++;
	}
	for(int i = 1; i <= n; i++) {
		if(du[i] < 3) continue;
		memset(c, 0, sizeof c);
		memset(c2, 0, sizeof c2);
		for(node *j = head[i]; j; j = j->nxt) {
			for(int k = 1; k <= maxdep; k++) t[k] = 0;
			maxdep = 0;
			dfs(j->to, i, 1);
			for(int k = 1; k <= maxdep; k++) {
				ans += (LL)c2[k] * t[k];
				c2[k] += t[k] * c[k];
				c[k] += t[k];
			}
		}
	}
	printf("%lld\n", ans);
	return 0;
}