[백준 4386] 별자리 만들기 - Python(파이썬)

우선순위 큐인 heapq 사용!

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):
    home[b] = a


n = int(sys.stdin.readline().strip())
arr = [list(map(float, sys.stdin.readline().split())) for _ in range(n)]
home = [-1] * (n + 1)

my_heap = []
answer = 0
cnt = 0

for i in range(n):
    for j in range(i+1, n):
        w, a, b = round(pow(pow(arr[i][0] - arr[j][0], 2) + pow(arr[i][1] - arr[j][1], 2), 0.5), 2), i, j
        heappush(my_heap, [w, a, b])

for i in range(len(my_heap)):
	# 우선순위 큐인 heapq를 사용, 최솟값 뽑기
    w, a, b = heappop(my_heap)
    f_a, f_b = find(a), find(b)
    # 해당 인덱스들 a와 b가 연결X일 경우,
    if f_a != f_b:
    	# 연결해주고 가중치 저장
        union(f_a, f_b)
        answer += w
        cnt += 1
		# n개의 정점을 가진 그래프의 최소 간선은 n-1개
        if cnt == n-1:
            break

print(answer)