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

문제 설명

문제에 사진이 많아 글로 작성하겠습니다.

루트 노드에서 순회를 시작했을 때, 처음으로 자식 노드가 2개 이상인 노드를 기가 노드라고 명칭한다. 기둥-가지를 줄여 기가 노드라 이름 붙였다.

여기서 루트 노드부터 기가 노드까지의 거리와 기가 노드로부터 가장 먼 리프 노드까지의 거리를 출력하는 문제이다.

입출력 및 제한 사항

풀이

2단계로 나누어서 풀이를 생각했다.

1. 루트노드가 아닌 노드들은 기가노드 이전까지 연결되어있는 모든 노드의 개수가 2개이므로 해당 조건으로 기가노드를 찾는다.

2. 기가노드로부터 DFS로 탐색하여 가장 멀리있는 노드까지의 거리를 구함.

루트노드로부터 직접 DFS를 탐색해도 괜찮지만 그렇게 된다면 기가 노드로부터의 거리를 다시 계산해야하기 때문에 DFS의 출발점을 기가노드로 잡았다.

내 코드

import sys
input = sys.stdin.readline
mis = lambda: map(int, input().split())

n, r = mis()
tree = [[] for _ in range(n+1)]
for _ in range(n-1):
    a, b, d = mis()
    
    tree[a].append((b, d))
    tree[b].append((a, d))
    
parent = 0
giga = r
giga_dis = 0

while True:
    deg = len(tree[giga])
    
    if giga == r:
        if deg != 1:
            break
    else:
        if deg != 2:
            break
    
    for nx, w in tree[giga]:
        if nx != parent:
            giga_dis += w
            parent, giga = giga, nx
            break

leaf_dis = 0
stack = [(giga, parent, 0)]
while stack:
    node, pa, dist = stack.pop()
    is_leaf = True
    
    for nx, w in tree[node]:
        if nx == pa:
            continue
        
        is_leaf = False
        stack.append((nx, node, dist + w))
        
    if is_leaf:
        if dist > leaf_dis:
            leaf_dis = dist

print(giga_dis, leaf_dis)

풀이는 평범한 DFS 풀이지만 재귀를 사용하지 않고 stack을 사용하여 풀이했다. 그 이유는 트리가 한 줄로 길게 이어질 수 있기 때문에 재귀 스택을 고려하여 stack을 선택했다. 신경써서 구현한 점은 방향 그래프라는 명시가 없었기 때문에 양방향이라고 가정하고 부모 노드를 다시 탐색하지 않는 것을 신경써서 풀이했다.

 

또한 기가 노드를 찾을 때 아래 코드에서 break를 해주지 않았다가 시간초과를 받았다...

for nx, w in tree[giga]:
    if nx != parent:
        giga_dis += w
        parent, giga = giga, nx
        break

처음엔 별 생각 없이 노드의 수가 2개니까 모두 순회했는데 처음에 자식노드를 탐색하면 parent가 바뀌기 때문에 이상한 루프에 빠질 수 있다는 것을 간과했다.