[BOJ-BFS] 1240 노드사이의 거리 Python

https://www.acmicpc.net/problem/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
from collections import deque

nodes, n = map(int, input().split())
# graph에 저장되는 값 (연결된 노드, 거리)
graph = [[] for _ in range(nodes+1)]

# graph에 노드간 연결, 거리 정보 저장
for _ in range(nodes-1):
    a, b, dist = map(int, input().split())
    graph[a].append((b, dist))
    graph[b].append((a, dist))

for _ in range(n):
    # visited로 거리 계산
    visited = [-1 for _ in range(nodes+1)]
    a, b = map(int, input().split())
    queue = deque()
    queue.append(a)
    visited[a] = 0
    while queue:
        now = queue.popleft()
        if now == b:
            print(visited[b])
            break

        for node, dist in graph[now]:
            if visited[node] == -1:
                # 거리 누적 (이전노드까지 걸린 거리 + 이전 노드에서 현재 노드까지 걸린 거리)
                visited[node] += (visited[now]+(dist+1))
                queue.append(node)

후기 🤔

처음 dfs로 시도했다가 계속 실패해서 bfs로 풀었다.
거리를 누적해주는 부분에서 조금 고민했는데, '최단거리'를 구하는 문제와 유사하다.
이렇게 '거리'를 구할 때는 visited 자료형을 활용하면 된다는걸 기억하자 !