백준 온라인 저지, DFS / 1240번: 노드사이의거리 (파이썬 / 백준 골드문제)

문제 주소: 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

힌트

접근 방법

- 리스트를 통해 트리를 구현한 뒤 DFS를 통해 거리를 측정한다.
- 트리로 구현됬기에 각 노드간의 거리는 DFS를 통해 이동했을 때가 가장 최소한의 거리이기에 찾고자하는 위치가 나온 경우 바로 리턴한 뒤, 값을 출력한다.

코드

# https://www.acmicpc.net/problem/1240
# 접근 방법
# 리스트를 통해 트리를 구현한 뒤 DFS를 통해 거리를 측정한다.
# 트리로 구현됬기에 각 노드간의 거리는 DFS를 통해 이동했을 때가 가장 최소한의 거리이기에 찾고자하는 위치가 나온 경우 바로 리턴한 뒤, 값을 출력한다.
def dfs(node1, node2, dist):
    visited[node1] = True
    if node1 == node2:
        return dist
    for x, d in nodes[node1]:
        if visited[x]:
            continue
        result = dfs(x, node2, dist+d)
        if result:
            return result
    return 0
    

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

n, m = map(int, input().split())
nodes = [[] for _ in range(n+1)]
for _ in range(n-1):
    u, v, cost = map(int, input().split())
    nodes[u].append([v, cost])
    nodes[v].append([u, cost])
for _ in range(m):
    u, v = map(int, input().split())
    visited = [False for _ in range(n+1)]
    print(dfs(u, v, 0))