AtCoder Beginner Contest 296E

ABC296E

言い換え
functional グラフが与えられる． 連結成分毎の サイクルの長さの和を求めよ．

実装
DFS する． 各頂点に対して，3つの状態がある．

• 0: 訪れてない
• 1: 訪れたが，まだ遷移の途中の点．
• 2: 訪れて，遷移が全て終わった点

気持ちとしては， 1 が計算中，2 が計算終了．

実装の注意 サイクルの長さが 1 or 2 の場合．

使っている記号，マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

int main() {

  ll n;

  cin >> n;

  vll a(n);

  vvll to(n);

  vll degree(n);

  rep(i,n) {

    cin >> a[i]; a[i]--;

    to[i].push_back(a[i]);

    degree[a[i]]++;

  }

  vll val(n);

  vll dep(n, -1);

  ll ans = 0;

  auto dfs = [&](auto f, ll cu, ll pa = -1, ll d = 0) -> void {

    if(val[cu] == 1){

      ans += d - dep[cu];

      return;

    }

    if(val[cu] == 2) return;

    val[cu] = 1;

    dep[cu] = d;

    for(auto ne: to[cu]) {

      f(f, ne, cu, d + 1);

    }

    val[cu] = 2;

  };

  rep(i,n) if(val[i] != 2){

    dfs(dfs,i);

  }

  cout << ans << endl;

  return 0;

}