[백준/BOJ] gold5 - 1240번 노드사이의 거리 (Python)

▶1240 - 노드사이의 거리

▶문제

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

 

▶입력

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

 

▶출력

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

728x90

▶풀이

일단 bfs를 이용해서 원하는 노드 간 거리를 구해줬다.

 

첫 번째 for문에서는 graph를 만들었고, 두 번째 for문에서는 두 노드를 입력받았다.

두 노드의 값을 bfs로 보내서 x에서 시작해 y에 도착하면 return해주는 방식을 사용했다.

 

골드 5 문제라고 나오는데, 사실 bfs를 안다면 쉽게 풀리는 문제라, 골드라고 쫄 필요 없어 보인다.

from collections import deque


def bfs(x, y):
    queue = deque()
    queue.append((x, 0))

    visit = [0 for _ in range(n + 1)]
    visit[x] = 1

    while queue:
        cur, cur_d = queue.popleft()
        if cur == y:
            return cur_d

        for nei, dist in arr[cur]:
            if visit[nei] == 0:
                visit[nei] = 1
                queue.append((nei, cur_d + dist))


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

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

for _ in range(m):
    a, b = map(int, input().split())
    answer.append(bfs(a, b))

for ans in answer:
    print(ans)