問題概要

頂点数 $N$ の functional graph が与えられる (出次数 1 の有向グラフ)。

このグラフの 2 頂点 $u, v$ からなる順序対 $(u, v)$ であって、頂点 $u$ から頂点 $v$ へと至るウォークが存在するものの個数を求めよ ( $u = v$ もあり)。

制約

$1 \le N \le 2 \times 10^{5}$

考えたこと

どの頂点からも、1 本の辺が出ているようなグラフを Functional Graph と呼ぶ。Functional Graph は「どの連結成分にもちょうど 1 個の閉路を含む」という特徴がある。

そして、下図のように、閉路に含まれないような頂点はすべて、閉路に向かって吸い込まれるような構造になっている。

drken1215.hatenablog.com

今回は、次の値を求めることにした。

dp[v] ← 頂点 $v$ から出発して到達できる頂点の個数

まず、Functional Graph のサイクルを検出して、検出されたサイクル (サイズを $m$ とする) の各頂点 $v$ について、dp[v] の値を $m$ としておく。

次に、グラフの各頂点 $v$ に対して、動的計画法 (メモ化再帰が楽) を用いて dp[v] の値を順次求めていく。

最後に、各頂点 $v$ についての dp[v] の値を合算すればよい。

全体の計算量は $O(N)$ となる。

コード

Functional Graph の閉路検出ライブラリを用いた。

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

// Edge Class
template<class T> struct Edge {
    int from, to;
    T val;
    Edge() : from(-1), to(-1), val(-1) { }
    Edge(int f, int t, T v = -1) : from(f), to(t), val(v) {}
    friend ostream& operator << (ostream& s, const Edge& E) {
        return s << E.from << "->" << E.to;
    }
};

// G[v] := 頂点 v から出ている辺
template<class T> struct CycleDetection {
    // input
    vector<Edge<T>> G;
    
    // intermediate results
    vector<bool> seen, finished;
    vector<int> history;
    
    // constructor
    CycleDetection() { }
    CycleDetection(const vector<Edge<T>> &graph) { init(graph); }
    void init(const vector<Edge<T>> &graph) {
        G = graph;
        seen.assign(G.size(), false);
        finished.assign(G.size(), false);
    }
    
    // return the vertex where cycle is detected
    int search(int v) {
        do {
            seen[v] = true;
            history.push_back(v);
            v = G[v].to;
            if (finished[v]) {
                v = -1;
                break;
            }
        } while (!seen[v]);
        pop_history();
        return v;
    }
    
    // pop history
    void pop_history() {
         while (!history.empty()) {
             int v = history.back();
             finished[v] = true;
             history.pop_back();
         }
     }
     
     // reconstruct
     vector<Edge<T>> reconstruct(int pos) {
         // reconstruct the cycle
         vector<Edge<T>> cycle;
         int v = pos;
         do {
             cycle.push_back(G[v]);
             v = G[v].to;
         } while (v != pos);
         return cycle;
     }
     
     // find cycle, v is the start vertex
     vector<Edge<T>> detect_from_v(int v) {
         int pos = search(v);
         if (pos != -1) return reconstruct(pos);
         else return vector<Edge<T>>();
     }
     
     // find all cycle
     vector<vector<Edge<T>>> detect_all() {
         vector<vector<Edge<T>>> res;
         for (int v = 0; v < (int)G.size(); ++v) {
             if (finished[v]) continue;
             int pos = search(v);
             if (pos == -1) continue;
             const vector<Edge<T>> &cycle = reconstruct(pos);
             if (!cycle.empty()) res.push_back(cycle);
         }
         return res;
     }
 };


int main() {
    long long N;
    cin >> N;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        --A[i];
    }
    
    // Functional Graph 構築
    vector<Edge<long long>> G(N);
    for (int i = 0; i < N; ++i) G[i] = Edge<long long>(i, A[i], 1);
    
    // 閉路検出
    using Cycle = vector<Edge<long long>>;
    CycleDetection<long long> cd(G);
    const vector<Cycle> &cycles = cd.detect_all();
    
    vector<long long> dp(N, -1);
    for (auto cycle : cycles) {
        for (auto e : cycle) dp[e.to] = cycle.size();
    }
    
    auto rec = [&](auto rec, int v) -> long long {
        if (dp[v] != -1) return dp[v];
        else return dp[v] = rec(rec, A[v]) + 1;
    };
    for (int v = 0; v < N; ++v) {
        if (dp[v] != -1) continue;
        rec(rec, v);
    }
    
    long long res = 0;
    for (int v = 0; v < N; ++v) {
        //cout << v << ": " << dp[v] << endl;
        res += dp[v];
    }
    cout << res << endl;
}