[Python/Gold5/1240] 노드사이의 거리

|    문제

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

|    입력

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

|    출력

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

 

 

---

정답 코드

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

class Solution:
    def dist_nodes(self):
        n, m = map(int, input().split())
        tree = defaultdict(list)
        for _ in range(n-1):
            v, k, w = map(int, input().split())
            tree[v].append((k,w))
            tree[k].append((v,w))
            
        find = [tuple(map(int, input().split())) for _ in range(m)]

        def dfs(node, dist, k):
            if node == dist:
                print(k)
                return
            
            for (v,w) in tree[node]:
                if not visited[v]:
                    visited[v] = True
                    dfs(v, dist, k+w)

        for (start, dist) in find:
            visited = [False] * (n+1)
            visited[start] = True
            dfs(start, dist, 0)


if __name__ == "__main__":
    s = Solution()
    s.dist_nodes()

 

 

|    비고

dfs/bfs 둘 다 가능한데, dfs로 풀어 보았다. 골드5 치곤 쉬운 문제인 것 같다.