[Python/파이썬] 백준 20924번 트리의 기둥과 가지

20924번: 트리의 기둥과 가지

 

문제

시청 공무원 마이크로는 과장으로부터 시에 있는 나무의 기둥의 길이와 가장 긴 가지의 길이를 파악하라는 업무 지시를 받았다.

마이크로는 ICPC Sinchon Winter Algorithm Camp에서 배운 트리 자료 구조를 이용하면 이 작업을 좀 더 수월하게 할 수 있으리라 판단했다. 

마이크로는 트리의 기둥과 가지를 분류하기 위해 기가 노드를 추가로 정의하였다.

기가 노드는 루트 노드에서 순회를 시작했을 때, 처음으로 자식 노드가 2$2$개 이상인 노드다. 기둥-가지를 줄여 기가 노드라 이름 붙였다. 위 그림에서 기가 노드는 $4$번 노드다.

단, 위 그림과 같이 리프 노드가 단 $1$개인 경우 리프 노드가 동시에 기가 노드가 된다.

또한, 위 그림과 같이 루트 노드가 동시에 기가 노드인 경우도 가능하다.

• 트리의 기둥은 루트 노드에서부터 기가 노드까지다. 위 그림에서 기둥은 $1-2-3-4$ 이다.
  기둥의 길이는 기둥의 간선 길이의 합인 $1 + 2 + 3 = 6$ 이 된다.
• 트리의 가지는 기가 노드에서부터 임의의 리프 노드까지다. 위 그림에서 가지는 $4-5-6-7$, $4-5-8$, 4−9$4-9$, $4-10-11$, $4-10-12$ 총 5$5$개가 있다.
  가지의 길이는 가지의 간선 길이의 합이다. 다행히도 가장 긴 가지의 길이 하나만 기재하면 된다. 4−10−12$4-10-12$ 가지가 간선 길이의 합 3+3=6$3 + 3 = 6$ 으로 가장 긴 가지이다.

마이크로는 시의 나무를 트리 자료 구조로 옮겼다. 그런데 과장이 마이크로에게 또 다른 업무를 지시했다! 너무 바쁜 마이크로를 대신해 트리의 기둥과 가장 긴 가지의 길이를 측정하자.

입력

첫 번째 줄에는 노드의 개수 $N$($1 \le N \le 200\,000$)과 루트 노드의 번호 $R$($1 \le R \le N$)이 주어진다.

이후 $N-1$개의 줄에 세 개의 정수 $a$, $b$, $d$($1 \le a, b \le N$, $ a \ne b$)가 주어진다. 이는 $a$번 노드와 $b$번 노드가 연결되어있으며 이 간선의 길이가 $d$($1 \le d \le 1\,000$)임을 의미한다. 노드는 1$1$번부터 $N$번까지 정수 번호가 매겨져 있으며 같은 간선은 여러 번 주어지지 않는다. 

트리가 아닌 그래프는 입력으로 주어지지 않는다.

 

출력

나무의 기둥의 길이와 가장 긴 가지의 길이를 출력한다.

코드

from collections import deque
import sys
input = sys.stdin.readline

n, r = map(int, input().split())
graph = [[] for _ in range(n+1)]
trunk = -1
longest_branch = 0
for _ in range(n-1):
    a, b, d = map(int, input().split())
    graph[a].append((b, d))
    graph[b].append((a, d))
q = deque([(r, 0)])
visited = [False] * (n+1)
visited[r] = True

# 기둥 길이 구하기
while q:
    now, dist = q.popleft()
    cnt = 0
    for nx, d in graph[now]:
        if not visited[nx]:
            q.append((nx, dist+d))
            visited[nx] = True
            cnt += 1
    if cnt >= 2 and trunk == -1:    # 기가 노드
        trunk = dist
    if cnt == 0:  # 리프 노드
        if trunk == -1:  # 리프 노드 == 기가 노드인 경우
            trunk = dist
        else:
            longest_branch = max(longest_branch, dist-trunk)

print(trunk, longest_branch)

BFS를 이용하는데 중간에 break하는 조건문이 없어서 트리의 모든 노드를 탐색하게 된다.

기가 노드인 노드는 탐색이 가능한 인접 노드, 즉 자식 노드가 2개 이상인 노드를 찾으면 된다. 리프 노드는 탐색이 가능한 노드가 없는 노드이다. 만약 기가 노드를 찾기 전에 리프 노드를 찾는다면 그 노드가 리프 노드인 동시에 기가 노드인 경우이므로 이 때는 루트 노드에서 현재 노드까지 탐색한 거리(dist)가 나무 기둥의 길이(trunk)이다.