[Python] 백준 #4386 별자리 만들기

📚문제: #4386 별자리 만들기(Gold 3)

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

• 별자리를 이루는 선은 서로 다른 두 별을 일직선으로 이은 형태이다.
• 모든 별들은 별자리 위의 선을 통해 서로 직/간접적으로 이어져 있어야 한다.
  별들이 2차원 평면 위에 놓여 있다. 선을 하나 이을 때마다 두 별 사이의 거리만큼의 비용이 든다고 할 때, 별자리를 만드는 최소 비용을 구하시오.

입력

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

출력

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

예제 입력 1

3
1.0 1.0
2.0 2.0
2.0 4.0

예제 출력 1

3.41

아이디어 및 구현

• 모든 별에 대해서 다른 모든 이웃 별과 거리를 계산하여 간선으로 저장한다. 간선은 총 n(n-1)/2개이며 양방향이다. 이후 Prim's Algorithm을 적용하여 최소 비용을 계산한다.

import sys
import math

# 별의 좌표가 주어집니다.
for i in range(n):
	
    # 소수점 이하 계산을 위해 'float'으로 입력받는다.
    x, y = map(float, sys.stdin.readline().split())

    stars.append([i, x, y])

# 모든 별에 대해서 이웃 별과의 거리를 구해줍니다.
for i in range(n):

    index1, x1, y1 = stars[i]

    for j in range(n):

        # 꺼낸 index가 j와 같을 때는 자기자신과의 거리계산이므로 skip
        if index1 == j:

            continue

        else:

            index2, x2, y2 = stars[j]
            dist = math.sqrt((x1-x2)**2 + (y1-y2)**2)
			
            # index1 별에서 index2 별로 가는 거리는 dist
            graph[index1].append((index2, dist))

전체 코드

import sys
import math
from heapq import heappush, heappop

# Minimum Spanning Tree : Prim Algorithm
def prim(start):

    # 총 거리
    total = 0

    # 간선을 선택할 때 활용할 우선순위 큐
    q = []

    # MST
    mst = []

    # 시작노드 처리
    mst.append(start)

    # 시작노드의 이웃을 큐에 추가 (거리, 별의 번호)
    for neighbor in graph[start]:

        heappush(q, (neighbor[1], neighbor[0]))
    
    while q:
        
        dist, cur = heappop(q)

        # mst에 이미 있다면 continue
        if cur in mst:

            continue
        
        # mst에 없다면 간선을 추가
        else:

            # mst에 해당 노드 추가
            mst.append(cur)

            # 총 거리에 선택 간선의 비용 추가
            total += dist

            # 이웃노드를 살펴봅니다.
            for neighbor in graph[cur]:

                # MST에 아직 없다면 큐에 추가해줍니다.
                if neighbor[0] not in mst:

                    heappush(q, (neighbor[1], neighbor[0]))

    # 답안 출력
    print(total)

# 별의 개수 n이 주어집니다.
n = int(sys.stdin.readline().rstrip())

# 그래프 초기화 (index : 0 ~ n-1)
graph =[[] for _ in range(n)]

# 별 리스트
stars = []

# 별의 좌표가 주어집니다.
for i in range(n):

    x, y = map(float, sys.stdin.readline().split())

    stars.append([i, x, y])

# 모든 별에 대해서 이웃 별과의 거리를 구해줍니다.
for i in range(n):

    index1, x1, y1 = stars[i]

    for j in range(n):

        # 꺼낸 index가 j와 같을 때는 자기자신과의 거리계산이므로 skip
        if index1 == j:

            continue

        else:

            index2, x2, y2 = stars[j]
            dist = math.sqrt((x1-x2)**2 + (y1-y2)**2)

            graph[index1].append((index2, dist))

# Prim Algorithm Execute
prim(0) # 어떤 별에서 시작해도 상관없다.