[백준][Python] 4386 별자리 그리기

문제

백준 4386 별자리 만들기

풀이

모든 별이 이어져 있는 최단 경로를 구하는 문제이므로 Prim 알고리즘을 사용하면 간단하게 풀 수가 있다. 단, 별과 별 사이의 거리를 구해서 직접 그래프를 구해줘야 하는 차이점이 있다.

전체 코드

from sys import maxsize
import heapq
import sys

def get_input():
    input = sys.stdin.readline
    N = int(input())
    star_list = [tuple(map(float, input().split())) for _ in range(N)]
    return N, star_list

def dist(A, B):
    a1, a2 = A
    b1, b2 = B
    return ( (a1 - b1) ** 2 + (a2 - b2) ** 2 ) ** (1/2)

def make_adjency_matrix(N, star_list):
    adjency_matrix = [[0] * N for _ in range(N)]

    for ii in range(N):
        for jj in range(ii+1, N):
            cost = dist(star_list[ii], star_list[jj])
            adjency_matrix[ii][jj] = cost
            adjency_matrix[jj][ii] = cost
    return adjency_matrix

def prim():
    visited = [False] * N
    heap = []
    total_cost = 0

    ## 0으로 시작하자 ~
    visited[0] = True
    for ii in range(1, N):
        heapq.heappush( heap, (adjency_matrix[0][ii], ii) )

    while(heap):
        cost, star = heapq.heappop(heap)
        if visited[star]:
            continue
        visited[star] = True
        total_cost += cost
        for ii in range(1, N):
            if visited[ii]:
                continue
            heapq.heappush( heap, (adjency_matrix[star][ii], ii))
    print(total_cost)

## graph만들고 prim하면 되지 않을까
if __name__ == "__main__":
    N, star_list = get_input()

    adjency_matrix = make_adjency_matrix(N, star_list)

    prim()