[Python/파이썬] 백준 1595번 북쪽나라의 도로

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

코드

import sys
input = sys.stdin.readline

MAX_NODE = 10_001

graph = [[] for _ in range(MAX_NODE)]
while True:
    try:
        a, b, dist = map(int, input().split())
        graph[a].append((b, dist))
        graph[b].append((a, dist))
    except:
        break


def dfs(start):
    visited = [False] * MAX_NODE
    visited[start] = True

    stack = [(start, 0)]
    farthest_node = start
    farthest_dist = 0

    while stack:
        now, now_dist = stack.pop()

        if now_dist > farthest_dist:
            farthest_dist = now_dist
            farthest_node = now

        for nx, nx_dist in graph[now]:
            if not visited[nx]:
                visited[nx] = True
                stack.append((nx, now_dist + nx_dist))

    return farthest_node, farthest_dist


start = 1
for i in range(1, 10_001):
    if graph[i]:
        start = i
        break

farthest_node, _ = dfs(start)
_, farthest_dist = dfs(farthest_node)

print(farthest_dist)

 

문제의 조건을 보면 "특정 도시를 두 번 이상 지나가지 않고서 임의의 두 도시간을 이동하는 경로가 유일하다"라고 되어 있다. 이 말은 각 도시들 간의 도로를 그래프로 나타내보면 트리 형태라는 말이다. 

 

어떤 트리에서 가장 거리가 먼 두 정점의 거리, 즉 트리의 지름을 찾을 때는 다음과 같은 과정으로 찾을 수 있다.

 

1. 아무 정점이나 하나 찍어서 그 정점에서부터 가장 거리가 먼 정점을 찾는다.

2. 1번에서 찾은 정점에서부터 가장 거리가 먼 정점까지의 거리를 구하면 답이다.

 

위 코드에서 `dfs()` 함수는 인자로 주어진 start 노드에서 가장 거리가 먼 노드와 그 거리를 리턴한다. 이 함수를 이용하여 임의의 점에서 가장 먼 노드를 구하고, 그 노드에서 `dfs()`를 한 번 더 사용하여 트리의 지름을 구하면 된다.