【题解】P3565 [POI2014]HOT-Hotels

题目链接(洛谷)

题目大意

给定一棵大小为\(n\)的树,在树上选\(3\)个点,要求两两距离相等,求方案数。

\(n\le 5000\),内存限制:\(62.5 \operatorname{Mb}\)(大概开\(1.5\times 10^7\)\(\text{int}\)

分析

很有用的性质:树上三个点距离两两相等,深度较深的那两个点的\(\text{lca}\)到三点距离相等。

我们只需要将这个\(\text{lca}\)提上来当成根,再统计深度相同的点。先枚举这个根,再枚举每棵子树,因为要保证被选的点\(\text{lca}\)是我们枚举的树根,每个子树中至多选一个点。

定义\(f_i\)表示在已经遍历的子树的深度为\(i\)的点中选出两个点的合法方案数。\(g_i\)表示选出一个点的合法方案数。\(num_i\)表示当前子树中深度为\(i\)的点的数量。则有如下转移:

$$ \[\begin{aligned} ans&+=f_i \times num_i\\ f_i&+=g_i \times num_i\\ g_i&+=num_i \end{aligned}\]

$$

实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#include <bits/stdc++.h>

#define ll long long
using namespace std;

const int N = 5005;
int n;
ll num[N], f[N], g[N], maxd, ans;

struct node {
int to, nxt;
} e[N << 1];
int head[N], tmp;

void add(int x, int y) {
e[++tmp].to = y;
e[tmp].nxt = head[x];
head[x] = tmp;
}

void dfs(int x, int fa, ll dep) {
maxd = max(maxd, dep);
num[dep]++;
for (int i = head[x]; i; i = e[i].nxt) {
int v = e[i].to;
if (v == fa) continue;
dfs(v, x, dep + 1);
}
}

void sol() {
for (int i = 1; i <= n; i++) { //枚举每棵子树
memset(g, 0, sizeof(g)); //换根记得清空dp数组
memset(f, 0, sizeof(f));
for (int j = head[i]; j; j = e[j].nxt) {
int v = e[j].to;
maxd = 0;
memset(num, 0, sizeof(num));
dfs(v, i, 1); // dfs 处理出每个深度上的结点个数
for (int k = 1; k <= maxd; k++) { // 转移,maxd是最深的深度
ans += f[k] * num[k];
f[k] += g[k] * num[k];
g[k] += num[k];
}
}
}
}

int main() {
ios::sync_with_stdio(false);
cin >> n;
for (int i = 1, x, y; i < n; i++) {
cin >> x >> y;
add(x, y);
add(y, x);
}
sol();
cout << ans << endl;
return 0;
}