[Python/파이썬] 백준 11085번 군사 이동

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

코드

from heapq import heappop, heappush
import sys
input = sys.stdin.readline


p, w = map(int, input().split())
c, v = map(int, input().split())

parent = [i for i in range(p)]


def find_parent(x):
    if x != parent[x]:
        parent[x] = find_parent(parent[x])
    return parent[x]


def union(a, b):
    a = find_parent(a)
    b = find_parent(b)

    if a == b:
        return

    if a < b:
        parent[b] = a
    else:
        parent[a] = b


roads = []
for _ in range(w):
    start, end, width = map(int, input().split())
    heappush(roads, (-width, start, end))

while roads:
    width, start, end = heappop(roads)
    union(start, end)
    if find_parent(c) == find_parent(v):
        print(-width)
        break

 

풀이 방법을 떠올리기가 어려웠을 뿐 기본적인 union-find 알고리즘을 안다면 쉽게 풀 수 있는 문제였다.

 

풀이 과정

1. 입력받은 길을 길이가 넓은 길이 앞에 오도록 `(-width, start, end)` 형태로 heap에 넣는다.

2. heap에서 하나씩 pop해서 union을 한다. while문은 c와 v의 루트가 같아진다면, 즉 같은 집합에 속하게 된다면 길이 연결된 것이므로 멈춘다. 멈추는 시점에 heappop된 길의 넓이가 만들어진 집합에서 가장 좁은 길이므로 이걸 출력하면 된다.