[백준] 1240 파이썬

문제

---

• 백준 1240 (골드 5)

백준 1240

가중치가 있는 트리에서 두 노드 사이의 거리를 출력하는 문제였다.
bfs 혹은 dfs로 비교적 간단하게 구현 가능해 보였다.

내 코드

---

bfs를 통해서 구현해 보았다.
시작 노드부터 bfs를 돌면서 타겟 노드가 나오면 그 거리를 리턴하여 출력할 수 있도록 구현하였다.

import sys
from collections import deque
input = sys.stdin.readline

n, m = map(int, input().split())
graph = [[] for _ in range(n + 1)]

for _ in range(n - 1):
    a, b, c = map(int, input().split())
    graph[a].append((b, c))
    graph[b].append((a, c))

target_list = [list(map(int, input().split())) for _ in range(m)]

def bfs(st, en):
    queue = deque([st])
    visited = [-1] * (n + 1)
    visited[st] = 0
    while queue:
        temp = queue.popleft()
        if temp == en:
            return visited[temp]
        for a, b in graph[temp]:
            if visited[a] == -1:
                visited[a] = visited[temp] + b
                queue.append(a)

for target in target_list:
    print(bfs(target[0], target[1]))