[백준1167번] 트리의 지름 / Python3

문제

트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

입력

트리가 입력으로 주어진다. 먼저 첫 번째 줄에서는 트리의 정점의 개수 V가 주어지고 (2≤V≤100,000)둘째 줄부터 V개의 줄에 걸쳐 간선의 정보가 다음과 같이 주어진다. (정점 번호는 1부터 V까지 매겨져 있다고 생각한다)

먼저 정점 번호가 주어지고, 이어서 연결된 간선의 정보를 의미하는 정수가 두 개씩 주어지는데, 하나는 정점번호, 다른 하나는 그 정점까지의 거리이다. 예를 들어 네 번째 줄의 경우 정점 3은 정점 1과 거리가 2인 간선으로 연결되어 있고, 정점 4와는 거리가 3인 간선으로 연결되어 있는 것을 보여준다. 각 줄의 마지막에는 -1이 입력으로 주어진다. 주어지는 거리는 모두 10,000 이하의 자연수이다.

출력

첫째 줄에 트리의 지름을 출력한다.

예제 입력 1

5
1 3 2 -1
2 4 4 -1
3 1 2 4 3 -1
4 2 4 3 3 5 6 -1
5 4 6 -1

예제 출력 1

11

나의 풀이

처음에 문제를 접했을 때 모든 노드에서 BFS로 탐색을 하고 가장 긴 거리를 반환하면 되겠구나라고 생각을 했지만, 문제의 입력조건에서 n이 10,000보다 작거나 같다했기때문에 시간복잡도때문에 해결할 수 없는 알고리즘이었다. 생각을 해보다 풀이를 찾아보니 트리의 특성을 이용해 풀어야했다.

"한 점에서 가장 먼 점을 A라하면 A에서의 가장 멀리 있는 점까지의 거리가 곧 트리의 지름이다."

위의 말이 곧 힌트이다. 처음에 글만 읽고는 바로 이해하지 못하였지만, 조금 생각해보니 깨달을 수 있었다. 좀 자세히 설명해보면 다음과 같다. 트리의 특성으로 어느 한 점을 잡고 들면 양쪽에 매달려오게 된다. 그림으로 설명해보면 다음과 같다.

예를들어, 다음과 같이 노드들이 연결되어있다고 해보자.
 1 - 2 - 3 - 4 - 5 

그럼 여기서 1번을 잡고 들어올리면 다음과 같다,
1
2
3
4
5

다음으로 2번을 들어올리면 다음과 같다.
   2
 1   3
       4
         5

마지막으로 예를들어 3번을 잡아 들어올리면 다음과 같다.
     3
   2   4
 1       5

이러한 특징을 생각하고 위의 말을 다시보자. 임의로 1번을 잡고 가장 먼 점을 찾아보자. 5번이다. 그럼 5번에서 가장 먼 점은? 바로 1번이다. 이 예는 조금 극단적인 예이고, 다음으로 2번을 잡은 일반적인 예로 살펴보자. 2번을 잡고 가장 먼점을 찾으면 5번이 나오게 된다. 그리고나서 5번에서 가장 먼 점을 찾게 되면 1번을 찾을 수 있게된다.

이와 같은 특성으로 모든 노드를 반복해 최장거리를 찾는 것이 아니라 두번의 BFS만으로 최장거리를 찾을 수 있게된다.

코드

# 1167번 트리의 지름
from collections import deque

def bfs(position):
    global max_dis, max_pos, visited

    q = deque()
    q.append([position,0])
    visited[pos] = True
    while q:
        cur, dis = q.pop()

        visited[cur] = True
        for n_node, d in adj_list[cur]:
            if not visited[n_node]:
                visited[n_node] = True
                q.append([n_node, dis + d])

                if dis + d > max_dis:
                    max_dis = dis + d
                    max_pos = n_node


n = int(input())
visited = [False] * (n+1)
max_dis = 0
max_pos = 0

adj_list = [[] for _ in range(n+1)]
for _ in range(n):
    tmp = list(map(int, input().split()))

    node = tmp.pop(0)
    while not tmp[0] == -1:
        n_node = tmp.pop(0)
        d = tmp.pop(0)

        adj_list[node].append([n_node, d])

# 1번에서 가장 먼 점을 찾음        
bfs(1)

# 다시 BFS를 돌리기위한 초기화
visited = [False] * (n+1)
max_dis = 0

# 위에서 찾은 가장 먼 점에서 다시 한번 먼 점을 찾음
bfs(max_pos)

print(max_dis)