【题解】P3565 [POI2014]HOT-Hotels

发表于 分类于 阅读次数： Valine：

题目链接（洛谷）

题目大意

给定一棵大小为$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$的点的数量。则有如下转移：

实现

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;
}