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

· · 题解

思路

观察 FakeLCA 函数代码。

int FakeLCA(int u, int v) {
    while (u != v) u = fa[u], v = fa[v];
    return u;
}

发现就是两个点不断向上跳的过程，思考 (u,v) 有贡献的情况：

这四种情况有重叠，考虑去重，首先统计 1 和 4 两种情况。

ans+=2*(n-1)+n;

接着 dfs 一遍处理出每棵子树的大小和每个点的深度，统计 2 和 3 两种情况。

void dfs(int u){
    siz[u]=1;dep[u]=dep[fa[u]]+1;
    for(auto v : G[u]){
        if(v==fa[u]) continue;
        dfs(v);
        siz[u]+=siz[v];
    }
}
void calc(int u){
    ++cnt[dep[u]];
    for(auto v : G[u]){
        calc(v);
    }
}
//main函数中
dfs(rt);
for(auto v : G[rt]){
    ans+=siz[v]*(n-1-siz[v]);// u 和 v 在不同子树
    //已经统计了(rt,v)与(v,rt)的贡献所以要"-1"
    for(int i=1;i<=n;i++) cnt[i]=0;
    calc(v);
    for(int i=1;i<=n;i++) ans+=cnt[i]*(cnt[i]-1);
}

在计算根节点子树贡献时已经包括了 u 和 v 不在同一颗子树而 dep[u]=dep[v] 的情况，所以要单独计算根节点每一棵子树内部各深度的出现次数记录在 cnt[] 中，计算下一棵子树时记得清空。

Code

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=1e6+5;

int fa[N],n,rt,siz[N],ans,dep[N],cnt[N];
vector<int> G[N];
void dfs(int u){
    siz[u]=1;dep[u]=dep[fa[u]]+1;
    for(auto v : G[u]){
        if(v==fa[u]) continue;
        dfs(v);
        siz[u]+=siz[v];
    }
}
void calc(int u){
    ++cnt[dep[u]];
    for(auto v : G[u]){
        calc(v);
    }
}
signed main(){
    cin>>n;
    for(int i=1;i<=n;i++){
        cin>>fa[i];
        if(fa[i]==i) rt=i;
        else G[fa[i]].push_back(i);
    }
    dfs(rt);
    ans+=2*(n-1)+n;
    for(auto v : G[rt]){
        ans+=siz[v]*(n-1-siz[v]);
        for(int i=1;i<=n;i++) cnt[i]=0;
        calc(v);
        for(int i=1;i<=n;i++) ans+=cnt[i]*(cnt[i]-1);
    }
    cout<<ans;
    return 0;
}