https://torus711.hatenablog.com/entry/2024/06/09/181410

問題文

　https://atcoder.jp/contests/abc357/tasks/abc357_e

問題概要

　$n$ 頂点の Functional Graph $G = ( V, E )$ が与えられる．具体的には，

$V = \{ 1, 2, \dots, n \}$
$|V| = |E| = n$
$A_i \in \{ 1, 2, \dots, n \}$ な $n$ 項の列 $A$ が与えられ，各 $i$ について $( i, A_i ) \in E$

である．
　頂点の順序対 $( u, v )$ であって，$0$ 本以上の辺を辿ることで $u$ から $v$ へ到達可能なものはいくつあるか？

制約

$1 \leq n \leq 2 \times 10^5$
$1 \leq A_i \leq n$

解法

　やや知識寄りかもですが重要な点として，Functional Graph の辺をたどっていくとサイクルに行き着くというのがあります．頂点 $u$ から到達可能な頂点というのは，$u$ がサイクルを構成する頂点であれば

$u$ を含むサイクルに含まれる全頂点

で，そうでないとき

$u$ 自身
頂点 $A_u$ から到達可能な頂点

です．なので，サイクルに含まれる頂点たちについてそのサイクルサイズを求めた後，それ以外の頂点について「適切な順序」で値を求めていけば，（その和をとることで）答えが求まります．サイクルの部分も含めて，「適切な順序」とは「トポロジカル順序の逆順」なので，何らかの方法でトポロジカル順序を求めて値を順に求めていけば答えを計算できます．
　持っていれば（あるいは AtCoder Library (ACL) を利用すれば）強連結成分分解を行うことでトポロジカル順序も求められるので，実装が楽になるかもしれません．
　強連結成分分解を利用する場合，計算量は（$|E| = n$ なことから）$\Theta( n )$ 時間となります．

コード

// 強連結成分分解 / Strongly Connected Component
// O( |V| + |E| )
class StronglyConnectedComponents; // 中身省略
// connect( from, to )
// solve() : 連結成分の数が返る
// componentNumbers() : 頂点番号→連結成分の番号 の表が返る
// 連結成分の番号の大小は、トポロジカル順序の大小と一致

int main()
{
	IN( int, N );
	VI A( N );
	cin >> A;
	MAP_PRED( A );

	StronglyConnectedComponents scc( N );
	REP( i, N )
	{
		scc.connect( i, A[i] );
	}

	const int K = scc.solve();
	const VI belongs_to = scc.componentNumbers();
	VI vertices( N );
	iota( ALL( vertices ), 0 );
	sort( ALL( vertices ), [&]( const int u, const int v ){
			return belongs_to[u] > belongs_to[v]; } );

	VI component_sizes( K );
	for_each( ALL( belongs_to ), [&]( const int c ){ ++component_sizes[c]; } );

	const auto cycle_size = [&]( const int u )
	{
		return component_sizes[ belongs_to[u] ];
	};
	const auto is_in_cycle = [&]( const int u )
	{
		return A[u] == u || 2 <= cycle_size( u );
	};

	VI res( N );
	FOR( u, vertices )
	{
		if ( is_in_cycle( u ) )
		{
			res[u] = cycle_size( u );
		}
		else
		{
			res[u] = 1 + res[ A[u] ];
		}
	}

	cout << accumulate( ALL( res ), 0LL ) << endl;

	return 0;
}