[ BOJ ] 4386 : 별자리 만들기 ( GOLD 3 ) / Python

>> 문제 바로가기 (https://www.acmicpc.net/problem/4386)

문제

도현이는 우주의 신이다. 이제 도현이는 아무렇게나 널브러져 있는 n개의 별들을 이어서 별자리를 하나 만들 것이다. 별자리의 조건은 다음과 같다.

• 별자리를 이루는 선은 서로 다른 두 별을 일직선으로 이은 형태이다.
• 모든 별들은 별자리 위의 선을 통해 서로 직/간접적으로 이어져 있어야 한다.

별들이 2차원 평면 위에 놓여 있다. 선을 하나 이을 때마다 두 별 사이의 거리만큼의 비용이 든다고 할 때, 별자리를 만드는 최소 비용을 구하시오.

입력

첫째 줄에 별의 개수 n이 주어진다. (1 ≤ n ≤ 100)

둘째 줄부터 n개의 줄에 걸쳐 각 별의 x, y좌표가 실수 형태로 주어지며, 최대 소수점 둘째자리까지 주어진다. 좌표는 1000을 넘지 않는 양의 실수이다.

출력

첫째 줄에 정답을 출력한다. 절대/상대 오차는 10-2까지 허용한다.

풀이 과정

모든 노드에 대한 점과 점 사이의 거리 가중치 간선을 만들고 MST 알고리즘을 수행한다.

import sys
import heapq
input = sys.stdin.readline

n = int(input().rstrip())
arr = [list(map(float, input().rstrip().split())) for _ in range(n)]
graph = [{} for _ in range(n+1)]

def Distances(ax, ay, bx, by):
    return ((ax - bx)**2 + (ay - by)**2)**(0.5)

for i in range(1, n+1):
    for j in range(i+1, n+1):
        graph[i][j] = Distances(arr[i-1][0], arr[i-1][1], arr[j-1][0], arr[j-1][1])
        graph[j][i] = Distances(arr[i-1][0], arr[i-1][1], arr[j-1][0], arr[j-1][1])
    
def prim():
    queue = []
    visit = [False] * (n + 1)
    sumw = 0
    heapq.heappush(queue, (0, 1))
    
    while queue:
        w, v = heapq.heappop(queue)
        if not visit[v]:
            visit[v] = True
            sumw += w
            for i in range(1, n+1):
                if graph[v].get(i) and not visit[i]:
                    heapq.heappush(queue, (graph[v][i], i))
    return sumw
               
print('%.2f'%prim())