본문 바로가기

백준 문제풀이 코드저장소/Gold

Baekjoon 1167. 트리의 지름 / Python

728x90

두 번의 DFS를 통해 문제를 해결했다.

임의의 노드를 정해서 가장 먼 노드로 움직인 후, 그 노드에서 다시 가장 먼 노드를 찾는다.

이 길이가 트리의 지름이 된다.

import sys
input = sys.stdin.readline

# 트리의 정점의 개수
v = int(input())
tree = [[] for _ in range(v + 1)]

# 간선의 정보를 입력받아 트리 구성
for _ in range(v):
    i, *arr = map(int, input().split())
    if len(arr) == 1:
        continue
    # 간선 정보를 트리에 추가
    for j in range(len(arr) // 2):
        tree[i].append((arr[2 * j], arr[2 * j + 1]))

# 거리 배열 초기화 (-1은 아직 방문하지 않았음을 의미)
dist = [-1] * (v + 1)
dist[1] = 0

# DFS 함수 정의
def dfs(s=1, tmp=0):
    # 현재 노드에서 인접한 노드들 탐색
    for node, distance in tree[s]:
        if dist[node] == -1:  # 아직 방문하지 않은 노드라면
            dist[node] = distance + tmp  # 거리 업데이트
            dfs(node, distance + tmp)  # 해당 노드에서 다시 DFS 수행

# 첫 번째 DFS 수행
dfs()

# 첫 번째 DFS에서 가장 먼 노드를 찾음
res = dist.index(max(dist))

# 거리 배열 초기화
dist = [-1] * (v + 1)
dist[res] = 0

# 두 번째 DFS 수행
dfs(res)

# 트리의 지름 출력
print(max(dist))
반응형