[BOJ/20924] 트리의 기둥과 가지, c++

www.acmicpc.net/problem/20924

20924번: 트리의 기둥과 가지

한줄 후기: 코드가 너무 더러운 것 같다..

 

//AC
//BOJ 20924 트리의 기둥과 가지

#include <iostream>
#include <vector>
using namespace std;
vector<pair<int, int> > v[200001];
int visit[200001];
int gi = 0, gidoong, ans = 0;
void dfs(int node, int pa, int R){
	visit[node] = 1;
	if(v[node].size() == 2){
		if(node == R){
			gidoong = node;
			return;
		}
		else{
			if(v[node][0].first != pa){
				gi += v[node][0].second;
				dfs(v[node][0].first, node, R);
			}
			else{
				gi += v[node][1].second;
				dfs(v[node][1].first, node, R);
			}
		}
	}
	else if(v[node].size() > 2){
		gidoong = node;
	}
	else if(v[node].size() == 1){
		if(node == R){
			gi += v[node][0].second;
			dfs(v[node][0].first, node, R);
		}
		else gidoong = node;
	}	
	return;
}
void maxCost(int node, int pa, int cur){
	visit[node] = 1;
	if(v[node].size() == 1){
		ans = max(ans, cur);
		return;
	}
	for(int i=0; i<v[node].size(); i++){
		pair<int, int> next = v[node][i];
		if(next.first == pa)	continue;
		if(!visit[next.first]){
			cur += next.second;
			maxCost(next.first, node, cur);
			cur -= next.second;
		}
	}
	
}
int main(){
	ios_base :: sync_with_stdio(false); 
	cin.tie(NULL); 
	cout.tie(NULL);
	int N, R;
	cin >> N >> R;
	for(int i=0; i<N-1; i++){
		int a, b, d;
		cin >> a >> b >> d;
		v[a].push_back({b, d});
		v[b].push_back({a, d});
	}
	if(N == 1){
		cout << gi << " " << gidoong;
		return 0;
	}
	dfs(R, 0, R);
	maxCost(gidoong, gidoong, 0);
	cout << gi << " " << ans;
	return 0;
}

기가노드를 먼저 구했다.

기가노드는 dfs 느낌으로 재귀를 통해 구하는데 만약 벡터에 자식이 3개 이상 있다면 (하나는 자식이 아니라 부모) 그 노드가 바로 기가노드가 된다.

2개 일때는 조건을 나눠야 하는데 만약 현재 노드가 루트라면 부모가 없으므로 2개라는건 바로 자식이 2개란 것

따라서 바로 기가노드가 된다.

그게 아니라면 cost를 더해주면서 재귀를 돌린다. 

만약 벡터에 연결된 노드가 한개 밖에 없다면, 현재 노드가 루트일 경우 재귀 , 아닐 경우 그 노드가 리프노드임으로 기가노드가 된다.

 

그 후, maxCost 함수를 통해 가장 긴 가지의 길이를 구한다. 

자식 노드들에게 재귀를 돌리는데 재귀를 돌고 나오면 다시 부모로 돌아오는 거니 cur에서 cost를 빼줘야 한다.

리프노드에 도착하면 현재 가지의 길이와 이전 max가지 길이와 비교해 더 큰 값을 저장해둔다.