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

백준 20924 - 트리의 기둥과 가지

시간 제한 2.5초(추가 시간 x), 메모리 제한 1024MB

# 조건

• 시청 공무원 마이크로는 과장으로부터 시에 있는 나무와 기둥의 길이와 가장 긴 가지의 길이를 파악하라는 업무 지시를 받았다.
• 마이크로는 ICPC Sinchon Winter Algorithm Camp에서 배운 트리 자료 구조를 이용하면 이 작업을 좀 더 수월하게 할 수 있으리라 판단했다.

• 마이크로는 트리의 기둥과 가지를 분류하기 위해 기가 노드를 추가로 정의하였다.
• 기가 노드는 루트 노드에서 순회를 시작했을 때, 처음으로 자식 노드가 $2$개 이상인 노드다.
• 기둥-가지를 줄여 기가 노드라 이름 붙였다. 위 그림에서 기가 노드는 $4$번 노드다.
  • 단, 리프 노드가 1개인 경우 리프 노드가 동시에 기가 노드가 되며, 루트 노드가 동시에 기가 노드인 경우도 가능하다.
• 마이크로는 시의 나무를 트리 자료 구조로 옮겼다.
• 그런데 과장이 마이크로에게 또 다른 업무를 지시했다!
• 너무 바쁜 마이크로를 대신해 트리의 기둥과 가장 긴 가지의 길이를 측정하자.

입력

• 첫 번재 줄에는 노드의 개수 N(1<=N<=200,000)과 루트 노드의 번호 R(1<=R<=N)이 주어진다.
• 이후 N-1개의 줄에 세 개의 정수 a, b, d(1<=a, b<=N, a != b)가 주어진다.
• 이는 a번 노드와 b번 노드가 연결되어있으며 이 간선의 길이가 d(1<=d<=1000)임을 의미한다.
• 노드는 1번부터 N번까지 정수 번호가 매겨져 있으며 같은 간선은 여러 번 주어지지 않는다.
• 트리가 아닌 그래프는 입력으로 주어지지 않는다.

출력

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

# 접근 방법

• 트리 구조를 활용하여 풀어준다.
  • 이 때, 주의할 점은 주어지는 간선의 정보가 부모 - 자식순서가 아니라는 것과, 루트 노드가 기가 노드일 수도 있다는 것을 인지하자.
• 주어지는 트리의 정보를 tree리스트에 양방향으로 기록해준다.
• 이후, visted 배열을 생성하고 우선적으로 find_giga함수를 실행해준다.
• q에 루트 노드를 넣은 후 다음 노드를 탐색할 건데 이 때 기가 노드인지 판단하는 조건문은
  • 기록되어 있는 길이가 3이상이거나, 길이가 2면서 루트 노드일 떄 기가 노드가 될 수 있다
  • 양방향으로 기록하므로 루트 노드는 기가 노드가 아니라면 1개의 정보만 존재해야 되므로 2개가 되는 경우 기가 노드이다.
  • 또한, 내부 노드의 경우 원래는 부모, 자식 총 2개의 정보가 있어야되지만 기가 노드의 경우 3개가 되므로 위와 같이 해주었다.
  • 또한, 위의 조건을 만족하지 못하여 return 없이 탐색이 끝났다면, 리프 노드가 기가 노드인 경우이므로 -1을 리턴해준다.
• 위의 함수로 받은 return이 만약 -1이라면 visted의 가장 큰 값과 0을 출력해주고
• 그렇지 않다면 max_length 함수를 실행해준다.
• 일반적인 bfs와 같이 다음 노드까지의 거리를 누적해서 q에 담아주고, max_val이 현재 누적된 거리보다 작다면 갱신해주면 된다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
from collections import deque  

def find_giga(st):  
    q = deque()  
    q.append(st)  
    while q:  
        node = q.popleft()  
        if len(tree[node]) >= 3 or (len(tree[node]) == 2 and node == st):  
            return node  
        for next_node, val in tree[node]:  
            if not visited[next_node] and next_node != st:  
                visited[next_node] = visited[node] + val  
                q.append(next_node)  
    return -1  

def max_length(st, R):  
    q = deque()  
    q.append((st, 0))  
    max_val = 0  
    while q:  
        node, val = q.popleft()  
        if val > max_val:  
            max_val = val  
        for next_node, next_val in tree[node]:  
            if not visited[next_node] and next_node != R:  
                visited[next_node] = 1  
                q.append((next_node, val + next_val))  
    return max_val  
N, R = map(int, input().split())  
tree = [[] for _ in range(N+1)]  
for _ in range(N-1):  
    a, b, c = map(int, input().split())  
    tree[a].append((b, c))  
    tree[b].append((a, c))  

visited = [0] * (N+1)  
giga = find_giga(R)  
if giga == -1:  
    print(max(visited), 0)  
else:  
    print(visited[giga], max_length(giga, R))