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

4386번: 별자리 만들기

문제

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

• 별자리를 이루는 선은 서로 다른 두 별을 일직선으로 이은 형태이다.
• 모든 별들은 별자리 위의 선을 통해 서로 직/간접적으로 이어져 있어야 한다.

별들이 2차원 평면 위에 놓여 있다. 선을 하나 이을 때마다 두 별 사이의 거리만큼의 비용이 든다고 할 때, 별자리를 만드는 최소 비용을 구하시오.

입력

첫째 줄에 별의 개수 n이 주어진다. (1 ≤ n ≤ 100)

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

출력

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

코드

전형적인 MST(Minumum Spanning Tree) 문제이다. 크루스칼 알고리즘과 프림 알고리즘 두 가지로 풀이 가능한데 나는 프림 알고리즘을 이용하여 풀이하였다.

• 프림 알고리즘

import sys
input = sys.stdin.readline

INF = int(1e9)

n = int(input())
star = []
for _ in range(n):
    x, y = map(float, input().split())
    star.append((x,y))
    
def get_dist(a,b):
    return ((a[0] - b[0])**2 + (a[1] - b[1]) ** 2) ** (1/2)

distance = [[INF for _ in range(n)] for _ in range(n)]
for i in range(n):
    for j in range(n):
        distance[i][j] = get_dist(star[i],star[j])

selected = set([0]) # MST에 포함된 별
unselected = set([i for i in range(1, n)])
ans = 0

while len(selected) != n:
    min_dist = 2000
    tmp = -1
    for i in selected:
        for j in unselected:
            if distance[i][j] < min_dist:
                tmp = j
                min_dist = distance[i][j]
    selected.add(tmp)
    unselected.remove(tmp)
    ans += min_dist
print(round(ans, 2))

1. 별 사이의 거리들을 미리 구해 distance 배열에 저장해둔다.
2. 임의의 한 별을 고르고 이 별을 MST 집합에 추가한다.
3. 이제 포함되지 않은 별들 중 MST에 포함된 별들과의 거리가 가장 가까운 별을 하나씩 선택하여 MST에 집어넣는 과정을 반복한다. 

• 크루스칼 알고리즘

import heapq
from itertools import combinations
import sys
input = sys.stdin.readline

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

parent = [i for i in range(n)]
def find_parent(x):
    if parent[x] != x:
        parent[x] = find_parent(parent[x])
    return parent[x]

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)
    
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

def get_dist(a, b):
    return ((a[0] - b[0])**2 + (a[1] - b[1]) ** 2) ** (1/2)

distance = []
for a, b in combinations(range(n), 2):
    heapq.heappush(distance, (get_dist(star[a], star[b]), a, b))

num = 0
ans = 0    
while num != n - 1:
    dist, a, b = heapq.heappop(distance)
    if find_parent(a) != find_parent(b):
        union_parent(a, b)
        num += 1
        ans += dist
        
print(round(ans, 2))

1. 별 사이의 거리를 모두 구하여 최소힙에 넣어 오름차순으로 정렬해준다. 이때 각 별의 번호(0 ~ n-1)도 2, 3번째 원소로 함께 넣어준다. n이 최대 100까지만 들어오기 때문에 모두 다 구해도 무리가 없다.
2. 노드의 개수가 n개인 최소 스패닝 트리는 간선의 개수가 (n - 1)개이다. (n - 1)개의 간선을 선택할 때까지 아래의 과정을 반복한다
   • 최소힙에서 원소를 하나씩 뽑아서 살펴본다. 간선에 속하는 두 노드의 루트 노드가 같지 않다면 사이클을 형성하지 않는 간선이므로 최소 스패닝 트리에 들어갈 수 있다. 뽑은 간선을 union_parent 함수로 합쳐주고 뽑은 간선의 개수(num)를 1 더해주고 우리가 구할 정답인 ans의 값에 이번에 뽑은 간선의 길이를 더해준다.