021.md

021 - Come Back in One Piece (★5)

解答

#include <iostream>
#include <vector>
#include <algorithm> // std::reverse(), std::fill()

using Graph = std::vector<std::vector<int>>;

void DFS(const Graph& g, int from, std::vector<bool>& visited, std::vector<int>& results)
{
	visited[from] = true;

	for (const auto& to : g[from])
	{
		if (!visited[to])
		{
			DFS(g, to, visited, results);
		}
	}

	// 帰りがけ順に頂点を記録
	results.push_back(from);
}

void DFS2(const Graph& g, int from, long long& count, std::vector<bool>& visited)
{
	visited[from] = true;

	++count;

	for (const auto& to : g[from])
	{
		if (!visited[to])
		{
			DFS2(g, to, count, visited);
		}
	}
}

int main()
{
	// N 頂点, M 辺
	int N, M;
	std::cin >> N >> M;

	// 入力されたグラフ
	Graph g1(N);
	
	// 辺の向きを反転したグラフ
	Graph g2(N);

	for (int i = 0; i < M; ++i)
	{
		int from, to;
		std::cin >> from >> to;
		--from;
		--to;

		g1[from].push_back(to);
		g2[to].push_back(from);
	}

	// 頂点が訪問済みか記録する配列
	std::vector<bool> visited(N);

	// 帰りがけ順に頂点を記録する配列
	std::vector<int> results;

	for (int i = 0; i < N; ++i)
	{
		if (!visited[i])
		{
			DFS(g1, i, visited, results);
		}
	}

	// 帰りがけ順が大きい頂点から処理するため並び順を反転
	std::reverse(results.begin(), results.end());

	// 次の探索のために、訪問済み頂点をリセット
	std::fill(visited.begin(), visited.end(), false);

	long long answer = 0;

	for (int from : results)
	{
		if (visited[from])
		{
			continue;
		}
		
		long long count = 0;

		// 行き止まりになるまでの頂点数を count に記録
		DFS2(g2, from, count, visited);

		// そのグループの頂点数から、answer に追加するペア数を求める
		answer += (count * (count - 1) / 2);
	}

	// 解答を出力
	std::cout << answer << '\n';
}