4386: 별자리 만들기

문제 링크

4386: 별자리 만들기

구현 방식

• Kruskal 알고리즘을 통해 MST를 구해주는 문제이다
• 별들이 2차원 좌표로 나타나기 때문에 edges 정의만 신경써서 해주면 된다
  -> 두 개의 별에 대해 nodes의 인덱스를 a, b라고 한다면, edges에 (math.sqrt((nodes[a][0]-nodes[b][0])**2 + (nodes[a][1]-nodes[b][1])**2), a, b)를 넣어줌. (cost, a, b) 형식

코드

import sys
import math

def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    a = find(a); b = find(b)
    if a < b: parent[b] = a
    else: parent[a] = b

N = int(sys.stdin.readline().strip())
nodes = []
for n in range(N): nodes.append(tuple(map(float, sys.stdin.readline().strip().split())))

edges = []
for i in range(N-1):
    for j in range(i+1, N):
        edges.append((math.sqrt((nodes[i][0]-nodes[j][0])**2 + (nodes[i][1]-nodes[j][1])**2), i, j))
edges.sort()

parent = [i for i in range(N+1)]
total_cost = 0
for cost, a, b in edges:
    if find(a) != find(b):
        union(a, b)
        total_cost += cost
print(round(total_cost, 2))