P11624 迷宫寻路 Round 3] 挂掉的模板 解题报告

P11624 迷宫寻路 Round 3] 挂掉的模板解题报告

前言：

挺好的一题，对我着这种老久没打过树上问题的人来说刚刚好。

思路：

首先，我们看 FakeLCA 函数实际上是 \(u,v\) 一个一个往上跳，跳到公共点的意思，特别的，树根 \(rt\) 的父节点是自己。那么对于一棵根节点的子树，以子树的根为 \(u\)，以子树根的子节点为 \(v\) 这种情况显然是不行的。我们考虑有哪些给出贡献的情况。

1. \((u,u)\):单个节点（样例1），贡献了 \(n\)。

2. \((rt,u) (u,rt)\):根可以等待下面的跳上来，贡献了 \(2\times(n-1)\)。

3. 两点连线经过 \(rt\) 的：只需选取任意两棵子树上的节点任意组合，我们枚举 \(rt\) 的子节点 \(u,v\)，以其为根的子树（显然包含根）大小为 sz[u],sz[v]，每组贡献了 sz[u]*sz[v] 。复杂度处在 \(O(n^2)\)，我们考虑优化到 \(O(n)\)：

   \[\begin{equation} \begin{aligned} &\sum_{u\in son(rt)}\sum_{v\in son(rt),u\neq v}sz_u*sz_v\\ &=\sum_{u\in son(rt)}sz_u*\sum_{v\in son(rt),u\neq v}sz_v\\ &=\sum_{u\in son(rt)}sz_u*(n-1-sz_u)\\ \end{aligned} \end{equation} \]

4. 两点连线不经过 \(rt\) 的：此时需要层数 dep[u]=dep[v]。我们分别对每一棵根的子树记录各层的节点数量，其中同一层的节点显然都是合法的，cnt[dep] 表示记录同层节点数，每一棵子树都有贡献 \(\sum cnt_{dep_i}\times (cnt_{dep_i}-1)\)，注意清空。

代码：

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int maxn=1e6+10;
int n,root,dm,ans=0,dep[maxn+10],cnt[maxn+10],sz[maxn+10];
vector<int>g[maxn+10];
void dfs(int u){
	sz[u]=1;
	cnt[dep[u]]++;
	for(auto v:g[u]){
		dep[v]=dep[u]+1;
		dm=max(dm,dep[v]);
		dfs(v);
		sz[u]+=sz[v];
	}
	if(dep[u]==1){
		for(int i=2;i<=dm;++i)ans+=cnt[i]*(cnt[i]-1),cnt[i]=0;
		dm=0;
	}
}
signed main(){
	ios::sync_with_stdio(0),cin.tie(0);
	cin>>n;
	for(int i=1;i<=n;++i){
		int f;
		cin>>f;
		if(f==i)root=i;
		else g[f].push_back(i);
	}
	dfs(root);
	for(auto i:g[root]){
		ans+=sz[i]*(n-1-sz[i]);
	}
	cout<<ans+n+(n-1)*2;
	
	return 0;
}