[골드 5] 백준 - 노드사이의 거리 (파이썬)

https://www.acmicpc.net/problem/1240

풀이

트리의 노드 N, 거리 구할 쌍의 개수 M을 초기에 입력받는디ㅏ.

N - 1개의 간선정보와 M개의 노드 쌍이 주어진다.

1. 트리는 사이클이 없고, 두 노드사이의 경로가 항상 유일하게 존재한다.
2. 따라서 s번 노드에서 DFS를 돌리면 모든 노드까지의 거리를 구할 수 있다.
3. t번 노드까지의 누적 거리를 출력하면 s - t 거리이다.

for _ in range(N - 1):
    u, v, w = map(int, input().split())
    u -= 1
    v -= 1
    adj[u].append((v, w))
    adj[v].append((u, w))

트리의 간선을 입력 받고, 편의를 위해서 1씩 빼준다.

또한 양방향으로 연결하기 u, v를 각각 삽입한다.

visit = [False] * N
dist = [0] * N

def dfs(u):
    visit[u] = True
    for v, w in adj[u]:
        if not visit[v]:
            dist[v] = dist[u] + w
            dfs(v)

DFS로 탐색하며 dist[v] = dist[u] + w 를 통해서 누적 거리를 계산한다.

visit 배열을 통해서 중복 방문을 방지한다.

for _ in range(M):
    s, t = map(int, input().split())
    s -= 1
    t -= 1

    visit = [False] * N
    dist = [0] * N

    dist[s] = 0
    dfs(s)

    print(dist[t])

각 거리 쿼리마다 s에서 출발해 DFS로 거리를 측정하고, dist[t]가 s에서 t까지의 거리이다.

코드

# 백준 1240 - 노드사이의 거리
# 분류 : 트리

N, M = map(int, input().split())
adj = [[] for _ in range(N)]

for _ in range(N - 1) :
    u, v, w = map(int, input().split())
    u -= 1
    v -= 1

    adj[u].append((v, w))
    adj[v].append((u, w))

visit = [False] * N
dist = [0] * N

def dfs(u) :
    visit[u] = True

    for v, w in adj[u] :
        if not visit[v] :
            dist[v] = dist[u] + w
            dfs(v)

for _ in range(M) :
    s, t = map(int, input().split())
    s -= 1
    t -= 1

    visit = [False] * N
    dist = [0] * N

    dist[s] = 0
    dfs(s)

    print(dist[t])