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

그래프 탐색 문제.

처음에 그냥 정점마다 DFS를 모두 돌려서 제출했는데 맞췄다. (시간 복잡도를 계산해도 가능할 것 같았다.)

그런데 문제가 의도하는 바가 이것도 아닌 것 같고, 제출 내역에 다른 사람들 시간을 보니 다시 풀어야 겠다고 생각해서

다시 풀어서 제대로 맞출 수 있었다.

DFS를 자유자재로 변형할 수 있으면 쉽게 맞출 수 있는 문제였다고 생각한다.

 

문제풀이

우선 처음 이 문제를 풀었을 때, 정점의 개수가 최대 10,000개인 것을 확인했고

간선의 개수가 최대 9,999개인 것을 확인했다.

각 정점마다 DFS를 한 번 하는데 소요되는 시간은 모든 정점을 방문했을 때, O(N)의 시간이 소요된다.

그러면 모든 정점에서 DFS를 한다면 총 소요되는 시간은 O(N^2)이다.

N은 최대 10,000이므로, 최악의 경우 100,000,000번의 연산을 하게 된다.

이는 1초이고, 문제의 제한시간은 2초이므로 시간 내에 풀 수 있게 된다.

 

/*
*/
#include <iostream>
#include <vector>
using namespace std;

int nodeCnt;
vector<pair<int, int>> adjList[10001];
bool visit[10001][10001];
int ans;

void dfs(int startNode, int curNode, int curDist) {
	visit[startNode][curNode] = true;
	ans = max(ans, curDist);

	for (pair<int, int> p : adjList[curNode]) {
		int nodeNum = p.first;
		int nodeDist = p.second;
		if (!visit[startNode][nodeNum]) {
			dfs(startNode, nodeNum, curDist + nodeDist);
		}
	}
}

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

	cin >> nodeCnt;

	int parNode, childNode, distance;
	for (int i = 1; i < nodeCnt; i++) {
		cin >> parNode >> childNode >> distance;
		adjList[parNode].push_back({ childNode, distance });
		adjList[childNode].push_back({ parNode, distance });
	}

	for (int i = 1; i <= nodeCnt; i++) {
		dfs(i, i, 0);
	}

	cout << ans;

	return 0;
}

위 코드가 O(N^2)의 수행을 하는 코드이다.

어려운 부분은 없고 다음과 같은 순서로 이루어진다.

• 노드의 개수를 입력받는다.
• 노드마다 연결된 노드를 입력받는다. (자식과 부모를 모두 저장해 놓는다.)
• 각 노드마다 dfs 메서드를 돌린다. 매개변수는 (시작 정점, 현재 탐색중인 정점, 현재 정점까지 거리)이다.
• 각 탐색하는 정점마다 ans와 현재 정점까지 거리를 비교해서 더 큰 값을 ans에 저장해놓는다.
• ans를 출력한다.

이렇게 하면 딱 1000ms 시간이 걸려서 통과할 수 있다. 그런데 이 문제와 똑같은 이름의 다른 문제가 있다.

그 문제는 골드 3인데, 정점 개수가 100,000개이다.

1167번: 트리의 지름

이 문제인데, 정점 개수가 100,000개이면 위 풀이로는 시간초과가 날 것이다.

 

따라서 다른 풀이를 찾아야 했다.

 

어떻게 하면 DFS 한 번으로 구할 수 있을까?

먼저 트리의 구조를 생각해보았다.

각 노드마다 연결된 자식들 중에서 가장 깊숙한 거리까지 뻗을 수 있는 두 개의 거리를 구한다.

두 거리를 합한 값이 그 노드에서 구할 수 있는 최대 지름이 된다.

말로 설명하려니 어렵다. 그림으로 살펴보자.

문제에서 보여주는 예시 그림이다.

[7] ~ [12] 노드는 자식이 없다. 가능한 지름은 모두 0이 된다.

[4]번 노드를 살펴보면, 최대로 깊숙하게 뻗을 수 있는 두 개의 거리는 1, 7이다. 따라서 [4]에서 구할 수 있는 최대 지름은 8이 된다.

[5]번 노드는 위와 마찬가지인 방법으로 살펴보면 가능한 지름은 19,

[6]번은 16이 된다.

[2]를 살펴보자. [2]에서는 한 개의 자식만 있다. 가장 깊숙하게 뻗을 수 있는 거리는 12이다. 다른 거리는 없다.

트리의 지름이 가능한 것은 12가 된다.

[3]을 보면, 자식이 두 개가 존재한다. 각각 뻗을 수 있는 최대 거리가 26, 19이다. 따라서 가능한 지름은 45이다.

[1]을 보면, 자식이 두 개가 존재한다. 각각 뻗을 수 있는 최대 거리가 15, 28이다. 따라서 가능한 지름은 43이다.

 

트리의 지름은 가장 긴 길이인 45가 된다.

 

이제 어떻게 구하는지는 알았다. 코드로는 어떻게 구현할 수 있을까?

나는 DFS를 통해서, 각 정점마다 부모에게 최대거리를 return하는 식으로 만들었다.

부모 노드들은 각 자식 노드들의 최대 거리만 필요하다.

 

따라서 부모 노드에서 각 자식 노드마다 가능한 거리는

[ 부모노드에서 자식1까지 거리 + 자식1의 최대한 깊숙하게 뻗은 거리, 부모노드에서 자식2까지 거리 + 자식2의 최대한 깊숙하게 뻗은 거리 ... 부모노드에서 자식k까지 거리 + 자식k의 최대한 깊숙하게 뻗은 거리 ]가 된다. 이 거리들 중 가장 긴 거리 2개를 더한 값이 해당 노드에서 가능한 트리의 지름이 되고,

가장 긴 값을 부모 노드에게 return 해주면 된다.

 

코드를 살펴보자.

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

int nodeCnt;
vector<pair<int, int>> adjList[10001];
vector<int> childDistanceVec[10001];
int ans;

사용한 STL과 변수

• nodeCnt : 입력받는 노드의 개수 N.
• adjList 벡터 : 인접한 노드의 정보를 입력받는 벡터. first에는 노드의 번호, second에는 해당 노드까지 거리를 담는다.
• childDistanceVec 벡터 : 한 정점에서 연결된 자식들부터 시작해서 뻗어나갈 수 있는 거리를 저장해두는 벡터. 입력받은 후 정렬시킨 뒤에 사용한다.
• ans : 출력하는 트리의 지름

 

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

	cin >> nodeCnt;

	int parNode, childNode, distance;
	for (int i = 1; i < nodeCnt; i++) {
		cin >> parNode >> childNode >> distance;
		adjList[parNode].push_back({ childNode, distance });
	}

	dfs(1);

	cout << ans;

	return 0;
}

메인함수

 

처음에 노드의 개수 N개를 입력받고,

입력받는 정보를 adjList 벡터에 넣어놓는다.

이 때, [부모 -> 자식]은 연결하지만 [자식 -> 부모]는 연결하지 않는다. 

왜냐하면 자식 -> 부모로 가는 길은 어차피 나중에 보여줄 dfs 메서드에서 return 하면서 최대 거리를 return 해주기 때문이다.

 

모든 정보를 입력받은 후에 dfs 메서드를 시작한다 매개변수인 1은 root Node 1번부터 시작한다는 의미이다.

 

dfs 메서드가 끝나면 ans를 출력한다.

 

int dfs(int nodeNum) {

	for (pair<int, int> p : adjList[nodeNum]) {
		int nextNodeNum = p.first;
		int nextNodeDistance = p.second;
		
		int childDistance = nextNodeDistance + dfs(nextNodeNum);
		childDistanceVec[nodeNum].push_back(childDistance);
	}

	sort(childDistanceVec[nodeNum].begin(), childDistanceVec[nodeNum].end());

	if (childDistanceVec[nodeNum].empty()) {
		return 0;
	}
	else if (childDistanceVec[nodeNum].size() == 1) {
		int bigIdx = childDistanceVec[nodeNum].size() - 1;
		ans = max(ans, childDistanceVec[nodeNum][bigIdx]);
		return childDistanceVec[nodeNum][bigIdx];
	}
	else {
		int firstBig = childDistanceVec[nodeNum].size() - 1;
		int secondBig = childDistanceVec[nodeNum].size() - 2;
		ans = max(ans, childDistanceVec[nodeNum][firstBig] + childDistanceVec[nodeNum][secondBig]);
		return childDistanceVec[nodeNum][firstBig];
	}

}

DFS 메서드.

시작하면 연결된 정점들(자식들)을 하나씩 살펴본다.

그리고 가능한 거리를 childDistance 변수에 저장한다. 이 때, 자식 노드의 거리는 다시 한번 dfs()메서드를 호출해서 return되도록 한다. 나중에 나오겠지만 이렇게 하면 해당 자식에서 뻗어나갈 수 있는 최대 거리가 return된다.

 

모든 자식 노드의 거리들을 구했으면 sort() 메서드를 통해 벡터를 정렬해준다. 그런 후에

• 만일 자식 노드의 거리 벡터가 비었다면, 즉 맨 하위 노드라면 0을 return한다.
• 만일 자식 노드의 거리 벡터의 사이즈가 1이라면, 즉 자식이 하나라면 그 거리와 현재 저장된 ans와 비교해서 더 큰 값을 ans에 저장한다. 그리고 나서 해당 거리를 return한다.
• 만일 자식 노드의 거리 벡터의 사이즈가 2 이상이라면, 즉 자식이 둘 이상이라면 첫번째로 큰 거리와 두번째로 큰 거리를 구한다. 그러고 난 후에, 두 거리를 더한 값과 현재 저장된 ans 값을 비교해서 더 큰 값을 ans에 저장한다. 그리고 가장 큰 거리를 return한다.

 

처음 맞았습니다. 는 맨 위 설명했던 풀이이다. 각 정점마다 DFS를 돌리는 방식이다. 시간이 엄청 오래 걸린다.

 

두번째 맞았습니다. 는 각 정점마다 visit 배열을 만들어서 푼 방식이다.

 

사실 visit 배열을 만들 필요가 없는게, 어차피 부모- 자식 관계는 하나 뿐이므로 각 정점은 한 번만 방문이 된다. 이 방문 기록을 지운 후에 제출한 게 세번째 맞았습니다. 이다.

 

GitHub - kimyunseok/cpp: my study-record

코드 전문은 위에서 확인 가능하다.