問題概要
Functional graph において，サイクル部分と tree 部分を求める．
解法
Cycle 部分と tree 部分を分けて処理する方がミスが少ない． Tree 部分は，元のグラフの反転を考えれば dp 出来る．
Cycle 部分は，cycle の長さを求める．これは DFS で，vis の値を 0~2 にして区別する事で可能． また，一つ cycle を見つけるだけなら，\(V\)回程度移動を繰り返せば cycle 内部の点に移動出来るので， そこから遷移を繰り返せば良い．
注意
Functional graph \(G\) の反転は，functional graph とは限らない． 実際，\(G\) においては複数の頂点から一つの頂点への辺が有り得る．Tree と cycle が交わる部分にあたる．
int main() {
ll n;
cin >> n;
vvll to(n);
vvll to_op(n);
rep(i,n){
ll a; cin >> a;
--a;
to[i].emplace_back(a);
to_op[a].emplace_back(i);
}
vll dp(n);
// cycle
{
vll vis(n);
vll path;
auto dfs = [&](auto dfs, ll cu) -> bool {
if(vis[cu] == 2) return false;
if(vis[cu] == 1){ // cycle
vll cycle;
cycle.emplace_back(cu);
while(path.back() != cu){
cycle.emplace_back(path.back());
path.pop_back();
}
ll cycle_length = cycle.size();
for(auto i: cycle){
dp[i] = cycle_length;
}
return true;
}
vis[cu] = 1;
path.push_back(cu);
for(auto ne: to[cu]){
if(dfs(dfs, ne)) {
vis[cu] = 2;
return true;
}
}
vis[cu] = 2;
return false;
};
rep(i,n) if(vis[i] != 2){
path = vll();
dfs(dfs, i);
}
}
// remain
{
auto dfs2 = [&](auto dfs2, ll cu) -> void {
for(auto ne: to_op[cu]) if(dp[ne] == 0){
dp[ne] = dp[cu] + 1;
dfs2(dfs2, ne);
}
};
rep(i,n) if(dp[i] != 0){
dfs2(dfs2, i);
}
}
ll ans = 0;
rep(i,n) ans += dp[i];
cout << ans << endl;
return 0;
}
