[BOJ] 1240번 노드사이의 거리 - 파이썬

📒 문제

 $N$개의 노드로 이루어진 트리가 주어지고 M개의 두 노드 쌍을 입력받을 때 두 노드 사이의 거리를 출력하라.

입력

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

출력

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

✏️ 풀이

from collections import deque

N, M = map(int, input().split())
tree = [[] for _ in range(N+1)]

def bfs(a, b):
    q = deque()
    q.append((a, 0))
    visit = [0 for _ in range(N+1)]
    visit[a] = 1

    while q:
        cur, cur_d = q.popleft()
        if cur == b:
            return cur_d
        for nxt, dist in tree[cur]:
            if visit[nxt] == 0:
                visit[nxt] = 1
                q.append((nxt, cur_d + dist))

for i in range(N-1):
    a, b, d = map(int, input().split())
    tree[a].append((b, d))
    tree[b].append((a, d))

want = []
for j in range(M):
    a, b = map(int, input().split())
    want.append(bfs(a, b))

for w in want:
    print(w)```