[Python] 백준, 1967(트리)

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

1967번: 트리의 지름

개요

트리 알고리즘 문항이긴 하나, 해결법은 BFS다. 트리를 정리하는 자료구조로 보통 리스트나 딕셔너리가 사용되는 건 알고 있었는데, 문제별로 그 구조를 조정할 필요는 있을 듯하다. 사실 처음 문제 풀 때 꽤 어렵게 접근해서 쩔쩔매긴 했는데, BFS를 두번 진행한다는 식의 발상은 못했다

아이디어

- BFS를 이용하여 루트 노드(1)부터 탐색을 진행한다. 가장 큰 가중치를 갖는 노드를 찾아 그 노드의 인덱스를 return.

- 발견한 노드로부터 다시 한번 BFS 탐색을 진행하여 지름의 총길이를 구한다.

주의사항

- 자료구조를 정리할 때, 가장 큰 가중치를 갖는 노드로부터 올라오는 과정으로 진행되기 때문에 양방향으로 연결해야 한다.

import sys
from collections import deque

n = int(sys.stdin.readline())
tree = {i: [] for i in range(n + 1)}

for _ in range(n - 1):
    a, b, d = map(int, sys.stdin.readline().split())
    tree[a].append([b, d])
    tree[b].append([a, d])

'''
tree = [{} for _ in range(n + 1)]
for _ in range(n - 1):
    a, b, d = map(int, sys.stdin.readline().split())
    tree[a][b] = d
    tree[b][a] = d
'''


def bfs(x):
    q = deque()
    q.append([x, 0])  # 첫 노드는 가중치가 없다.
    visit = [False] * (n + 1)
    visit[x] = True
    result = [0, 0]

    while q:
        node, cnt = q.popleft()
        for i in tree[node]:
            a1, d1 = i[0], i[1]
            if not visit[a1]:
                visit[a1] = True
                q.append([a1, d1 + cnt])  # 아직 방문하지 않았다면 현재 노드, 이전 가중치에 현재 가중치 더해준 값을 append.
                if result[1] <= d1 + cnt:  # result[1]의 값과 현재 가중치까지의 합의 비교를 통해 result[1] 갱신
                    result[0] = a1
                    result[1] = d1 + cnt
    return result


print(bfs(bfs(1)[0])[1])