4386. 별자리 만들기 (Python)

4386번: 별자리 만들기

 

풀이

 모든 별들이 사이클이 없는 그래프로 연결되어 있으면 된다. 즉, N개의 별이 있을 때 N-1개의 간선을 사용하여 연결하는 최소 스패닝 트리를 구하면 된다. 

 

 최소 스패닝 트리를 만들기 위해서 사용되는 알고리즘에는 크루스칼 알고리즘과 프림 알고리즘이 있다. 일반적으로 크루스칼 알고리즘은 정점에 비해 간선의 개수가 적을 때 사용하면 좋고 프림 알고리즘은 정점에 비해 간선의 개수가 많을 때 사용하면 좋다. 해당 문제는 모든 별들끼리의 길이를 구한 후 최소 스패닝 트리를 완성해야 한다. 정점에 비해 간선의 개수가 많으므로 프림 알고리즘을 사용하면 조금 더 짧은 시간에 답을 구할 수 있다.

 

코드

import sys
from math import sqrt
input = sys.stdin.readline
INF = float('inf')


# 프림 알고리즘을 사용한 최소 스패닝 트리 구성 
# start : 임의의 시작 정점
def prim(start):
    distances[start] = 0 # 시작 정점의 거리는 0으로 설정 
    res = 0.0

    # 총 N번 반복하면 최소 스패닝 트리를 만들 수 있다.
    for _ in range(N):
        minStar, minDistance = -1, INF

        # 현재 그래프에서 도달할 수 있는 별 중에서 아직 연결되지 않았고 가장 가까운 거리에 있는 별을 찾는다.
        for i in range(N):
            if not visited[i] and distances[i] < minDistance:
                minStar, minDistance = i, distances[i]

        # 위에서 찾은 별에 대해서 방문 처리를 해주고 해당 길이를 결과에 더해준다.
        visited[minStar] = True
        res += distances[minStar]

        # 새롭게 연결된 별에서 이을 수 있는 간선들을 체크하여 distances 배열의 값을 업데이트한다.
        for nextStar, nextCost in costs[minStar]:
            if not visited[nextStar] and nextCost < distances[nextStar]:
                distances[nextStar] = nextCost

    return round(res, 2)


if __name__ == '__main__':
    N = int(input())
    stars = []
    costs = [[] for _ in range(N)]
    visited = [False for _ in range(N)]
    distances = [INF for _ in range(N)]

    for _ in range(N):
        stars.append(list(map(float, input().split())))

    # 모든 별들의 연결 길이를 구한다.
    for i in range(N):
        for j in range(i+1, N):
            x = abs(stars[i][0] - stars[j][0])
            y = abs(stars[i][1] - stars[j][1])
            cost = sqrt(x ** 2 + y ** 2)
            costs[i].append((j, cost))
            costs[j].append((i, cost))

    # 임의의 별 0번에서 시작 
    print(prim(0))