[백준 2887] 행성 터널 - Python(파이썬)

우선순위 큐인 heap 사용!

from heapq import heappop, heappush
import sys


def find(x):
    if home[x] < 0:
        return x
    home[x] = find(home[x])
    return home[x]


def union(a, b):
    a, b = find(a), find(b)
	
    # 빠른 탐색을 위해, 최솟값에 저장 
    if a != b:
        if home[a] < home[b]:
            home[a] += home[b]
            home[b] = a
        else:
            home[b] += home[a]
            home[a] = b


N = int(sys.stdin.readline().strip())
arr = [list(map(int, sys.stdin.readline().split())) + [i] for i in range(N)]
home = [-1] * N
my_heap = []
answer, cnt = 0, 0

for i in range(3):
    arr.sort(key=lambda x: x[i])
    for j in range(1, N):
        heappush(my_heap, (abs(arr[j][i] - arr[j-1][i]), arr[j][3], arr[j-1][3]))

for _ in range(len(my_heap)):
	# 우선순위 큐인 힙에서 최소값 가중치를 하나씩 뽑아
    w, a, b = heappop(my_heap)
	# 해당 좌표들인 a와 b가 연결X일 경우,
    if find(a) != find(b):
    	# 연결시켜주고, 가중치 저장
        union(a, b)
        answer += w
        cnt += 1
		# N개의 정점을 가진 그래프의 최소 간선 수는 N-1개
        if cnt == N-1:
            break

print(answer)