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

황승환·2022년 7월 6일
0

Python

목록 보기
344/498


이번 문제는 최소스패닝트리 문제였다. 처음 접해보는 알고리즘이었기 때문에 구현 방법을 찾아보고 구현하였다. 처음 접근 할 때에는 찾아보지 않고, permutations를 구하여 그 순서대로 방문하는 가중치 중 최솟값을 구하도록 구현하였다. 답은 잘 도출되었지만, 메모리초과가 발생하였다.

최소스패닝트리를 찾아보니 Kruskal과 Prim 알고리즘 중 하나를 선택할 수 있었다. 이번 문제는 Prim으로 접근해보기로 하였다. 우선 모든 점들 사이의 거리를 구하여 인접리스트 형식으로 저장하고, 최소힙을 이용하여 가중치가 작은 순으로 접근하게 한다. 방문처리를 진행하여 방문하지 않은 정점에 대하여 총 가중치를 계산하고, 최소힙에 넣는 방식을 n-1번 반복하면 결과를 얻을 수 있었다.

Code

import math
import heapq
n = int(input())
star = [tuple(map(float, input().split())) for _ in range(n)]
def get_dist(a, b):
    y1, x1 = a
    y2, x2 = b
    return round(math.sqrt((y1-y2)**2 + (x1-x2)**2), 2)
def prim(start, weight):
    visited = [False for _ in range(n)]
    q = [(weight, start)]
    answer = 0
    cnt = 0
    while cnt < n:
        w, cur = heapq.heappop(q)
        if visited[cur]:
            continue
        visited[cur] = True
        answer += w
        cnt += 1
        for nxt, nxt_w in link[cur]:
            if visited[nxt]:
                continue
            heapq.heappush(q, (nxt_w, nxt))
    return answer
link = [[] for _ in range(n)]
for i in range(n):
    for j in range(n):
        if i == j:
            continue
        dist = get_dist(star[i], star[j])
        link[i].append((j, dist))
print(prim(0, 0))

profile
꾸준함을 꿈꾸는 SW 전공 학부생의 개발 일기

0개의 댓글