[알고리즘/파이썬] 백준 4386 - 별자리 만들기

문제

https://www.acmicpc.net/problem/4386

4386번: 별자리 만들기

 

풀이

최소 스패닝 트리를 조금 활용한 문제...? 라고 보면 될 것 같다.

 

프림 알고리즘을 이용해서 풀었고, 일반적으로 그래프 문제들은 연결돼있는 간선과 노드의 정보를 주는 경우가 많은데 얘는 그 정보를 두 점 사이의 거리를 활용해서 직접 만들어야 하는 문제였다.

math.sqrt를 통해 제곱근을 구했고, 소수점 둘째자리까지만 활용하기 위해 round 함수를 사용했다.

 

이 외에는 일반적인 프림 알고리즘과 동일하게 작성했다.

 

 

코드

import math
import sys; input=sys.stdin.readline
import heapq


n = int(input())

stars = []
for _ in range(n):
    x,y = map(float,input().split())
    stars.append((x,y))
    
graph = [[] for _ in range(n)]

def prim(start,weight):
    visited = [0 for _ in range(n)]
    q = [[weight,start]]
    cnt = 0
    ans = 0

    while cnt < n:
        w,s = heapq.heappop(q)
        
        if visited[s] == 1: continue
        visited[s] = 1
        for next,c in graph[s]:
            heapq.heappush(q,(c,next))
        ans += w
        cnt += 1

    return ans

for i in range(n):
    for j in range(i+1,n):
        a = (stars[i][0] - stars[j][0]) ** 2
        b = (stars[i][1] - stars[j][1]) ** 2
        dis = round(math.sqrt(a+b),2)
        graph[i].append((j,dis))
        graph[j].append((i,dis))

result = prim(1,0)
print(result)