백준 20924 트리의 기둥과 가지

문제링크

https://www.acmicpc.net/problem/20924

문제

풀이

두 번의 dfs를 이용해 풀이했다.
1. 첫 번째 dfs에서 기둥의 길이와 기가 노드를 찾아주었다.
2. 두 번째 dfs의 시작점을 기가 노드로 지정하여 가장 긴 가지를 찾아주었다.

방문한 노드는 visit 배열을 통해 표시하여 두 번째 dfs 실행 시 기둥 노드에 방문하지 않도록 하였다.

코드

#include <iostream>
#include <vector>
#include  <algorithm>
using namespace std;

int n, r;
bool vi[200001];
vector<pair<int, int>> v[200001];

pair<int, int> find_trunk(int now, int cost) {

	vi[now] = true;
	if (v[now].size() > 2) return { cost,now };
	if (v[r].size()==2) return{ cost,now };

	for (int i = 0; i < v[now].size(); i++) {
		int next = v[now][i].second;
		int ncost = v[now][i].first;
		if (vi[next]) continue;
		return find_trunk(next, cost + ncost);
	}

	return { cost,now };
}

int find_leaf(int now, int cost) {

	vi[now] = true;
	if (v[now].size() == 1) return cost;

	int temp = 0;
	for (int i = 0; i < v[now].size(); i++) {
		int next = v[now][i].second;
		int ncost = v[now][i].first;
		if (vi[next]) continue;

		temp = max(temp, find_leaf(next, cost + ncost));
	}

	return temp;


}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	cin >> n >> r;
	for (int i = 0; i < n - 1; i++) {
		int a, b, d;
		cin >> a >> b >> d;
		v[a].push_back({ d,b });
		v[b].push_back({ d,a });
	}

	pair<int, int> ans1 = find_trunk(r, 0);
	int ans2 = find_leaf(ans1.second, 0);

	cout << ans1.first << ' ' << ans2;
}

후기

단순한 dfs 문제인데 예외 경우를 생각 못 해서 조금 헤맸다.
기둥 노드의 길이가 0이고 기가 노드에서 두 갈래로 뻗어 나가는 경우는 따로 처리를 해주었다.