백준 1967번 트리의 지름 C++

출처 : https://www.acmicpc.net/problem/1967

 

루트부터 시작해서 자식들 중 가장 긴 두 경로를 재귀적으로 탐색하는 방식으로 풀었다.

 

#include <iostream>
#include <string>
#include <vector>
#include <math.h>
#include <algorithm>
#include <utility>
#include <stack>
#include <queue>
#include <math.h>
#include <set>
#include <map>
#include <list>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <limits.h>

using namespace std;
using int64 = long long;

int ans = 0;

struct Node {
	vector<pair<int, int>> childs;
};

map<int, Node> mp;

int GetDiameter(int root) {

	if (mp.find(root) == mp.end())
		return 0;

	vector<int> dists(mp[root].childs.size());

	for (int i = 0; i < mp[root].childs.size(); i++) {
		dists[i] = mp[root].childs[i].second + GetDiameter(mp[root].childs[i].first);
	}

	sort(dists.begin(), dists.end(), greater<int>());

	int dist = dists[0];
	if (dists.size() >= 2)
		dist += dists[1];

	ans = max(ans, dist);

	return dists[0];
}

int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);

	int n;
	cin >> n;
	for (int i = 0; i < n - 1; i++) {
		int num, child, dist;
		cin >> num >> child >> dist;
		auto it = mp.find(num);
		if (it == mp.end()) {
			mp.insert({ num,Node() });
			mp[num].childs.push_back({ child,dist });
		}
		else {
			it->second.childs.push_back({ child,dist });
		}
	}
	GetDiameter(1);
	cout << ans;
};