[BOJ/Python] 4386: 별자리 만들기

gold 3 / boj-4386: 별자리 만들기

문제 흐름

모든 별(노드)를 연결할 때 최소 비용을 구하는 문제이므로 MST를 적용할 수 있다.

두개의 별을 연결하는 데는 두 별 사이의 거리 만큼의 비용이 소모되고,
미리 연결된 별은 없으며,
모든 별이 이어질 수 있는 잠재 간선으로 연결되어 있으므로,
노드에 비해 간선이 많다고 할 수 있다.
따라서 Prim 알고리즘을 적용하는 것이 Kruskal보다 효율적일 것이라고 생각했다.

정석적인 Prim 알고리즘 구현을 통해 해결했고,
초기에 모든 두 쌍의 별에 대해 그 거리를 구해 graph에 저장하여 간선을 모두 표시하였다.

코드

# 별자리 만들기
# MST

import sys
from collections import defaultdict
import heapq
input = sys.stdin.readline

def prim(graph):
    pq = [(0, 0)]
    visited = [0 for _ in range(len(graph))]
    result = 0
    
    while pq:
        weight, x = heapq.heappop(pq)
        if visited[x]:
            continue
        visited[x] = 1
        result += weight
        
        for nx, w in graph[x]:
            if not visited[nx]:
                heapq.heappush(pq, (w, nx))
    return result


if __name__ == "__main__":
    n = int(input())
    star = []
    for _ in range(n):
        x, y = map(float, input().split())
        star.append((x, y))
    
    graph = defaultdict(list)
    for i in range(n):
        for j in range(n):
            if i == j:
                continue
            ax, ay = star[i][0], star[i][1]
            bx, by = star[j][0], star[j][1]
            distance = ((ax-bx)**2 + (ay-by)**2)**0.5
            graph[i].append((j, distance))
            graph[j].append((i, distance))
    
    answer = prim(graph)
    print(answer)