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

백준 1240 - 노드사이의 거리

시간 제한 2초, 메모리 제한 128MB

# 조건

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

입력

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

출력

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

제한

• $2≤N≤1,000$ 
• $1≤M≤1,000$ 
• 트리 상에 연결된 두 점과 거리는 $10,000$ 이하인 자연수이다.
• 트리 노드의 번호는 $1$부터 $N$까지 자연수이며, 두 노드가 같은 번호를 갖는 경우는 없다.

# 접근 방법

• bfs를 활용하여 그래프 탐색을 해주면 된다.
• 입력받는 노드 간의 관계를 양방향으로 넣어주고, bfs를 돌려주면 된다.
• 이 때, M개의 줄에 주어지는 두 노드를 탐색할 때마다 VISITED 배열을 초기화해주면 된다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
from collections import deque  

def bfs(st, en):  
    q = deque()  
    q.append((st, 0))  
    visited[st] = True  
    while q:  
        cur_node, dist = q.popleft()  
        if cur_node == en:  
            print(dist)  
            return  
        for node, val in tree[cur_node]:  
            if not visited[node]:  
                visited[node] = True  
                q.append((node, dist + val))  

N, M = map(int, input().split())  
tree = [[] for _ in range(N+1)]  
for _ in range(N-1):  
    a, b, c = map(int, input().split())  
    tree[a].append((b, c))  
    tree[b].append((a, c))  
for _ in range(M):  
    visited = [False] * (N+1)  
    a, b = map(int, input().split())  
    bfs(a, b)