[ 백준 / P5 / 파이썬 Python ] 정점들의 거리

링크 : https://www.acmicpc.net/problem/1761

1761번: 정점들의 거리

 

1. 코드

---

      

      

      

    
    

      <python />
    
from collections import defaultdict
from heapq import heappop,heappush

def LCS(a,b):
    while (da := depth[a]) != (db:=depth[b]):
        if da > db:
            a = parent[a]
            continue
        b = parent[b]
    while a != b:
        a,b = parent[a],parent[b]
    return a

def DIJKSTRA(root):
    parent = [x for x in range(n + 1)]
    dist = [MAX] * (n + 1)
    dist[root] = 0
    depth = [0] * (n + 1)
    heap = [(0, root)]
    while heap:
        w, node = heappop(heap)
        for nw, nextNode in graph[node]:
            if dist[nextNode] > w + nw:
                depth[nextNode] = depth[node] + 1
                dist[nextNode] = w + nw
                heappush(heap, (w + nw, nextNode))
                parent[nextNode] = node
    return dist,parent,depth

n = int(input())
MAX = float('inf')
graph = defaultdict(list)
for _ in range(n-1):
    a,b,c = map(int,input().split())
    graph[a].append((c,b))
    graph[b].append((c,a))

dist,parent,depth = DIJKSTRA(1)

m = int(input())
for _ in range(m):
    a,b = map(int,input().split())
    lcs = LCS(a,b)
    print(dist[a]+dist[b]-2*dist[lcs])

 

2. 사용 알고리즘

---

"최소 공통 조상(LCS) + 다익스트라"

 

3. 풀이

---

      

      

      

    
    

      <python />
    
'''
- 공식 : A까지의 거리 + B까지의 거리 - 2*최소공통조상까지의 거리
- 루트노드부터 나머지 노드까지의 최단거리 생성 -> 루트노드가 무엇인지 알아야 함
    - 아무 노드나 하나 잡고 루트노드로 잡아도 무관하지 않을까?
    - 시간복잡도에서 차이가 나긴 할 것 같음.
- 루트노드를 잡고 거기서부터 탐색해나가면서 부모노드를 갱신해가면 될 듯.
'''

위의 내용은 문제 풀이에 들어가기 앞서 내 생각을 주석으로 정리한 것이다.

 

트리 구조에서 두 노드 사이의 거리를 구하기 위해서는 두 노드가 연결된 부분을 찾아 최단경로로 이동해야 한다.

그 공식이 바로 위에 적힌 "A까지의 거리 + B까지의 거리 - 2*최소공통조상까지의 거리" 이 부분이다.

 

하지만 최소 공통 조상을 찾기 위해서는 트리의 부모-자식 관계를 알아야 하는데 이 문제의 입력에는 두 노드의 연결관계, 거리만 보여주고 어떤 노드가 부모인지 자식인지는 파악할 수 없었다.

따라서 아무 노드나 하나를 선택해 루트노드로 만들고 기준을 세웠다. 이 방법이 어떤 테스트케이스의 경우에서는 시간복잡도가 느려질 수도 있으나 루트노드를 정함으로써 모든 연결관계에서 부모-자식을 나눌 수 있기 때문에 이 방법을 선택했다.

 

루트노드를 정한 후에는 다익스트라 알고리즘을 사용하여 루트노드로부터 나머지 모든 노드로의 최단거리를 구했다.

      

      

      

    
    

      <python />
    
def DIJKSTRA(root):
    parent = [x for x in range(n + 1)]
    dist = [MAX] * (n + 1)
    dist[root] = 0
    depth = [0] * (n + 1)
    heap = [(0, root)]
    while heap:
        w, node = heappop(heap)
        for nw, nextNode in graph[node]:
            if dist[nextNode] > w + nw:
                depth[nextNode] = depth[node] + 1
                dist[nextNode] = w + nw
                heappush(heap, (w + nw, nextNode))
                parent[nextNode] = node
    return dist,parent,depth

기존의 다익스트라와 동작과정에서 차이는 없으나 최소공통조상 알고리즘에서 사용할 부모노드 배열과 깊이 배열을 탐색하며 함께 갱신해주었다.

위의 코드를 통해 루트노드로부터 모든 노드까지의 최단거리를 구하고 트리의 모든 부모, 자식 관계를 생성했다.

 

이후에는 임의의 두 노드 A,B에 대하여 두 노드의 최소 공통 조상을 구하는 LCA 알고리즘을 작성했다.

      

      

      

    
    

      <python />
    
def LCS(a,b):
    while (da := depth[a]) != (db:=depth[b]):
        if da > db:
            a = parent[a]
            continue
        b = parent[b]
    while a != b:
        a,b = parent[a],parent[b]
    return a

다익스트라를 구현하며 함께 생성한 depth와 parent배열을 사용하였다.

트리에서 두 노드의 깊이가 같아질 때까지 반복한 후, 두 노드의 깊이를 동일하게 유지하면서 최소 공통 조상을 향해 탐색해 나아갔다.

 

마지막으로 정답을 출력하는 부분이다.

      

      

      

    
    

      <python />
    
m = int(input())
for _ in range(m):
    a,b = map(int,input().split())
    lcs = LCS(a,b)
    print(dist[a]+dist[b]-2*dist[lcs])

맨 위에서 적은 공식처럼 매번 입력이 들어올 때마다 "A까지의 거리 + B까지의 거리 - 2*최소공통조상까지의 거리"을 코드로 구현하였다.