【题解】P11624 [迷宫寻路 Round 3] 挂掉的模板

· · 题解

首先，翻译一下 FakeLCA 函数，就是“每次让 u 和 v 向上爬升，直至 u = v，视 u 为 lca(u,v)”。

记树的根为 root，树上节点 u 的深度为 d_u（d_{root}=1）。要使数对 (u,v) 产生贡献，发现可以分为几种情况：

  1. 对于根节点的任意两颗不同子树中的各一个节点 u 和 v，(u,v) 和 (v,u) 可以产生贡献。

  2. 对于在根节点的同一颗子树的两个节点 u 和 v，若 d_u = d_v，则 (u,v) 和 (v,u) 可以产生贡献。

对于第一种情况，直观的想法是，统计出根节点的每一颗子树大小，记做 s_1,s_2,\dots,s_k，则贡献为 \sum\limits_{i=1}^k \sum\limits_{j=i+1}^k s_i \times s_j \times 2。但这样时间复杂度是 O(n ^ 2)。我们不妨转化下思路，设当前 dfs 到第 i 颗子树，则第 i 颗子树与之前的任一子树各选一个节点组合的方案为 s_i \times \sum\limits_{j=1}^{i-1} s_j \times 2。这里 \sum\limits_{j=1}^{i-1} 是动态变化的，用一个变量记录即可。

对于第二种情况，记 t_i 表示当前子树深度为 i 的节点数量，贡献为 \sum t_i \times t_{i-1}。记得在 dfs 到下一颗子树时清空 t_i。

对于第三种情况，贡献为 n。

对于第四种情况，可以将其归入第一种情况，将根节点视为第一颗子树的一部分，即开始时 sum = 1。

详见代码（记得开 long long）：

#include <bits/stdc++.h>
#define int long long
using namespace std;

const int N = 1e6 + 5;

int hd[N], cnt, d[N], t[N], fa[N], mx;
int n, root, sum = 1, ans;

struct Edge {
    int to, nx;
} e[N];

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

int dfs(int u) {
    d[u] = d[fa[u]] + 1; t[d[u]] ++;
    mx = max(mx, d[u]);
        // mx 表示当前子树的最大深度
    int now = 1;
        // now 表示当前子树的节点数
    for (int i = hd[u]; i; i = e[i].nx) {
        int v = e[i].to;
        now += dfs(v);
    }
    if (d[u] == 2) {
        ans += now * sum * 2;
        sum += now;
            // sum 表示之前 dfs 过的子树的节点数之和（包括根节点）
        for (int i = 2; i <= mx; i ++) {
            ans += t[i] * (t[i] - 1), t[i] = 0;
        }
        mx = 0;
    }
    return now;
}

signed main() {
    cin >> n;
    for (int i = 1; i <= n; i ++) {
        cin >> fa[i];
        if (i == fa[i]) root = i;
        else add(fa[i], i);
    }
    dfs(root);
    cout << ans + n;
    return 0;
}