백준 1967: 트리의 지름

문제 요약

트리에서 임의의 두 노드 사이의 경로의 길이를 트리의 지름이라고 한다. 정확히 정의하자면 트리에 존재하는 모든 경로들 중에서 가장 긴 것의 길이를 말한다.

입력으로 루트가 있는 트리를 가중치가 있는 간선들로 줄 때, 트리의 지름을 구해서 출력하는 프로그램을 작성하시오.

문제 분류

• 그래프 이론
• 그래프 탐색
• 트리
• DFS

문제 풀이

각 트리를 DFS하며 재귀호출 하면 되는데, 결국 두 노드를 지정하여 양쪽으로 쫙 땡겨서 긴 길이를 만들려면, 그 두 노드는 모두 리프노드여야 할 것이다.따라서 리프노드에서 자기 노드까지의 최대 길이를 반환하면 좋겠다고 생각했다. 결국 k번째 노드를 탐색할때 k번째의 노드가 루트 노드라고 하면, (각 자식노드들의 최대 길이) + (자신과 자식과의 가중치)의 최대값을 반환하면서도, 루트에서 리프로 통하는 최대값의 경로 중 두 개만을 선택하여 합한 뒤, 최대치를 갱신해주면 된다.

노드 연결상태는 multimap으로, 루트노드를 기준으로 탐색하는 서브트리의 가중치의 최대치는 우선순위 큐로 해결하였다. 우선순위 큐에 일단 전부 넣고, 2개 이하로만 남긴 뒤, 나머지를 전부 더하여 실제 결과값에 최대치로 갱신해주었다.

풀이 코드

#include <stdio.h>
#include <iostream>
#include <map>
#include <queue>

using namespace std;

multimap<int, pair<int, int>> mm;
int n, res;

int dfs(int idx) {
	int ret = 0, sum = 0, temp;
	auto rg = mm.equal_range(idx);
	priority_queue<int, vector<int>, greater<int>> v;
	if (rg.first == rg.second) return 0;
	for (auto& it = rg.first; it != rg.second; it++) {
		temp = dfs(it->second.first) + it->second.second;
		ret = max(ret, temp);
		v.push(temp);
	}
	while (v.size() > 2)  v.pop();
	while (!v.empty()) {
		sum += v.top();
		v.pop();
	}
	res = max(res, sum);
	return ret;
}
	
int main() {
	int in1, in2, in3;
	cin >> n;
	for (int i = 1; i < n; i++) {
		scanf("%d%d%d", &in1, &in2, &in3);
		mm.insert({ in1, {in2, in3} });
	}
	dfs(1);
	cout << res;
	return 0;
}