039.md

039 - Tree Distance (★5)

解答

#include <iostream>
#include <vector>
#include <algorithm>

void DFS(const std::vector<std::vector<int>>& graph, int from, int prev, std::vector<int>& dp)
{
	dp[from] = 1;

	for (const auto& to : graph[from])
	{
		if (to != prev)
		{
			DFS(graph, to, from, dp);

			// 帰りがけ順で dp を更新
			dp[from] += dp[to];
		}
	}
}

int main()
{
	int N;
	std::cin >> N;

	std::vector<int> A(N - 1), B(N - 1);
	std::vector<std::vector<int>> graph(N);
	for (int i = 0; i < (N - 1); ++i)
	{
		std::cin >> A[i] >> B[i];
		--A[i]; --B[i];
		graph[A[i]].push_back(B[i]);
		graph[B[i]].push_back(A[i]);
	}

	std::vector<bool> visited(N);

	// 頂点 0 を根としたうえで,
	// dp[i]: i を根とする部分木の頂点数
	std::vector<int> dp(N);

	DFS(graph, 0, -1, dp);

	long long answer = 0;

	// すべての辺について
	for (int i = 0; i < (N - 1); ++i)
	{
		// 注目してる辺で葉側のグループの頂点数
		const long long a = std::min(dp[A[i]], dp[B[i]]);

		// その反対側のグループの頂点数
		const long long b = (N - a);

		// 組み合わせの数を加算
		answer += (a * b);
	}

	std::cout << answer << '\n';
}