1167. 트리의 지름 (Python)

1167번: 트리의 지름

 

풀이

 트리의 지름을 구하려면 임의의 노드에서 가장 멀리 있는 노드 u를 구한다. 그 후 다시 노드 u에서 가장 멀리 있는 노드 v를 구한다. 이렇게 구한 노드 u와 노드 v의 길이가 트리의 지름이 된다.

 

 이에 대한 증명은 링크를 참고하자...

 

코드

from collections import deque
import sys 

input = sys.stdin.readline

# node에서 가장 멀리 떨어져 있는 노드를 구하기 위한 bfs 함수 
def bfs(node):
    visited = [False for _ in range(V+1)]
    visited[node] = True 

    q = deque()
    q.append((node, 0))

    maxCost = 0 # 가장 멀리 떨어져 있는 거리
    resNode = -1 # maxCost에 해당하는 노드

    while q:
        nowNode, nowCost = q.popleft()

        for (nextNode, cost) in graph[nowNode]:
            if visited[nextNode]:
                continue 
            
            visited[nextNode] = True
            nextCost = nowCost + cost 
            q.append((nextNode, nextCost))

            if nextCost > maxCost:
                maxCost = nextCost
                resNode = nextNode
    
    global res 
    res = max(res, maxCost)

    return resNode


if __name__ == '__main__':
    V = int(input())
    graph = [[] for _ in range(V+1)]

    for _ in range(V):
        edges = list(map(int, input().split()))
        edges.pop()
        node = edges.pop(0)

        for i in range(0, len(edges), 2):
            nextNode = edges[i]
            cost = edges[i+1]
            graph[node].append((nextNode, cost))

    res = 0 

    u = bfs(1) # 임의의 노드 1번 노드에서 가장 멀리 떨어져 있는 노드 u를 찾는다.
    v = bfs(u) # 노드 u에서 가장 멀리 떨어져 있는 노드 v를 찾는다.

    # 노드 u와 노드 v의 거리가 트리의 지름이 된다.
    print(res)