[골드 4] 백준 1967 - 트리의 지름 (파이썬)

https://www.acmicpc.net/problem/1967

풀이

트리의 지름은 다음의 두번의 DFS로 구할 수 있다.

1. 임의의 노드에서 가장 먼 노드를 찾는다 → who
2. 그 노드에서 다시 가장 먼 노드까지의 거리가 트리의 지름이다.

import sys
sys.setrecursionlimit(10 ** 9)

파이썬의 기본 재귀 깊이가 얕기 때문에 DFS에서 런타임 에러 방지를 위해서 깊이를 늘린다.

N = int(input())
adj = [[] for _ in range(N)]

노드 수 N입력 후, 입접 리스트 adj를 생성한다.

for _ in range(N - 1):
    u, v, w = map(int, input().split())
    u -= 1
    v -= 1
    adj[u].append((v, w))
    adj[v].append((u, w))

양방향 간선 정보를 입력 받고, u - 1, w - 1은 노드 번호를 0-indexed로 변환한다.

visit = [False] * N
dist = [0] * N

def dfs(u):
    visit[u] = True
    for v, w in adj[u]:
        if not visit[v]:
            dist[v] = dist[u] + w
            dfs(v)

dfs(0)

첫번째 DFS를 수행한다. 즉 0번 노드에서 DFS를 수행하여 가장 멀리 떨어진 노드를 찾기 위한 거리를 계산한다.

max_dist = 0
who = 0
for i in range(N):
    if max_dist < dist[i]:
        max_dist = dist[i]
        who = i

그 후 0번 노드로부터 가장 먼 노드 who를 구한다.

visit = [False] * N
dist = [0] * N
dfs(who)

두번째 DFS를 수행하여 트리의 지름을 구한다. 가장 먼 노드 who에서 다시 DFS를 시작한다.

max_dist = 0
for i in range(N):
    if max_dist < dist[i]:
        max_dist = dist[i]

print(max_dist)

최종적으로 who에서 가장 먼 노드까지의 거리를 출력한다, 이 거리가 즉 트리의 거리이다.

시간 복잡도 분석

• DFS 2번 → O(N)

코드

# 백준 1967 - 트리의 지름
# 분류 : 트리
import sys
sys.setrecursionlimit(10 ** 9)

N = int(input())
adj = [[] for _ in range(N)]

for _ in range(N - 1) :
    u, v, w = map(int, input().split())
    u -= 1
    v -= 1

    adj[u].append((v, w))
    adj[v].append((u, w))

visit = [False] * N
dist = [0] * N

def dfs(u) :
    visit[u] = True

    for v, w in adj[u] :
        if not visit[v] :
            dist[v] = dist[u] + w
            dfs(v)

dfs(0)

max_dist = 0
who = 0

for i in range(N) :
    if max_dist < dist[i] :
        max_dist = dist[i]
        who = i

visit = [False] * N
dist = [0] * N
dfs(who)

max_dist = 0
for i in range(N) :
    if max_dist < dist[i] :
        max_dist = dist[i]

print(max_dist)