[Baekjoon] 노드사이의 거리 / Python, 파이썬 / 1240

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

1240번: 노드사이의 거리

문제 설명

문제

N(2≤N≤1,000)개의 노드로 이루어진 트리가 주어지고 M(M≤1,000)개의 두 노드 쌍을 입력받을 때 두 노드 사이의 거리를 출력하라.

입력

첫째 줄에 노드의 개수 N이 입력되고 다음 N-1개의 줄에 트리 상에 연결된 두 점과 거리(10,000 이하의 정수)를 입력받는다. 그 다음 줄에는 거리를 알고 싶은 M개의 노드 쌍이 한 줄에 한 쌍씩 입력된다.

출력

M개의 줄에 차례대로 입력받은 두 노드 사이의 거리를 출력한다.

예제 입력 1

4 2
2 1 2
4 3 2
1 4 3
1 2
3 2

예제 출력 1

2
7

문제 풀이

import sys

input = sys.stdin.readline
sys.setrecursionlimit(10**6)

def dfs(x):
    for data in graph[x]:
        y = data[0]
        cost = data[1]
        if not visited[y]:
            visited[y] = True
            distance[y] = distance[x] + cost
            dfs(y)


N, M = map(int, input().split())

graph = [[] for _ in range(N + 1)]

for _ in range(N - 1):
    x, y, dst = map(int, input().split())
    graph[x].append((y, dst))
    graph[y].append((x, dst))

for _ in range(M):
    x, y = map(int, input().split())
    visited = [False for _ in range(N+1)]
    distance = [-1 for _ in range(N+1)]
    visited[x] = True
    distance[x] = 0
    dfs(x)
    print(distance[y])