[python] 백준 4386 : 별자리 만들기

문제 링크
https://www.acmicpc.net/problem/4386

접근

문제를 읽자마자 저번에 풀었던 1647번 : 도시 분할 계획 문제가 떠올랐다.

1647번: 도시 분할 계획

1647번 문제는 Kruskal's 알고리즘으로 간선을 중심으로 최소신장트리를 만들어갔다면, 이번 문제는 Prim's 알고리즘으로 정점 중심으로 최소신장트리를 만들어야겠다고 생각했다.

알고리즘은 간단하다.
1. 시작점 설정
2. 갈 수 있는 모든 점에 대해 거리 계산 (다익스트라처럼)
3. 그 중 가장 최소인 점으로 이동
4. 2번으로

이 문제의 경우 다른 그래프 문제와 다르게, 사실 모든 점에서 모든 점으로 갈 수 있는 완전그래프로 구성된 문제이다. 따라서 한 점이 열리게 되면 그 점은 현재 자신이 방문했었던 점을 제외하고 모든 점에 갈 수 있게 된다. 그래서 memo 라는 이중 리스트를 만들어 각 점마다 떨어진 거리를 저장하는 작업을 먼저 하였다.

for i in range(n):
    for j in range(n):
        memo[i][j] = math.sqrt((ps[i][0] - ps[j][0]) ** 2 + (ps[i][1] - ps[j][1]) ** 2)
        # 자기 자신과의 거리는 0

이제 프림 알고리즘을 사용할 차례이다.
복잡도를 줄이기 위해 최소힙을 사용하였다.
점과 점 사이의 거리가 곧 가중치이므로, 가중치가 적은 녀석부터 뽑고, 뽑은 점에 대해서 새로운 길이 열린 것이라 볼 수 있으므로 최소힙에 memo[cur][] 에 해당하는 녀석들을 다시 집어넣는 것을 반복한다.

sum = 0
visit = set()
mh = []
heapq.heappush(mh, (0, 0))

while len(visit) <= n - 1:
    cost, cur = heapq.heappop(mh)
    if cur in visit:
        continue
    visit.add(cur)
    sum += cost
    for to in range(n):
        if cur != to and to not in visit:
            heapq.heappush(mh, (memo[cur][to], to))

정답 코드

import sys
import math
import heapq

sys.setrecursionlimit(10 ** 8)
input = lambda: sys.stdin.readline().rstrip()

n = int(input())
ps = []
for i in range(n):
    x, y = map(float, input().split())
    ps.append((x, y))

memo = [[0 for _ in range(n)] for _ in range(n)]

for i in range(n):
    for j in range(n):
        memo[i][j] = math.sqrt((ps[i][0] - ps[j][0]) ** 2 + (ps[i][1] - ps[j][1]) ** 2)

sum = 0
visit = set()
mh = []
heapq.heappush(mh, (0, 0))

while len(visit) <= n - 1:
    cost, cur = heapq.heappop(mh)
    if cur in visit:
        continue
    visit.add(cur)
    sum += cost
    for to in range(n):
        if cur != to and to not in visit:
            heapq.heappush(mh, (memo[cur][to], to))

print(f"{sum:.2f}")