[백준 파이썬] 4386 별자리 만들기

백준 4386

모든 별들을 이어주는데 필요한 비용의 최솟값을 출력해야한다. 최소 비용을 출력하므로 최소 스패닝 트리를 구하면 된다. 최소가 되려면 모든 노드를 포함하지만 두 노드 사이에는 하나의, 가장 비용이 적은 간선만 있도록 하면 된다.
(최소 스패닝 트리 문제는 에지 리스트와 유니온 파인드로 해결한다.)
노드 사이의 비용은 유클리디안 거리로 반복문으로 구해주면 된다.

from decimal import Decimal

def find(p): # 부모 노드 찾기
    if p == parent[p]:
        return p
    else: 
        return find(parent[p])

def union(p, q): # 부모 노드가 다르다면 두 노드 병합, 부모 노드는 가장 작은 값으로
    p = find(p)
    q = find(q)
    assert p != q

    if p < q:
        parent[q] = p
    else:
        parent[p] = q

N = int(input())
coords = [list(map(Decimal, input().split())) for _ in range(N)]

# 두 별 사이의 유클리디안 거리를 구하고 이 세 값으로 에지 리스트를 만들어준다
edges = []
for i in range(N):
    for j in range(i+1, N):
        dist = ((coords[i][0] - coords[j][0])**2 
                + (coords[i][1] - coords[j][1])**2) ** Decimal('0.5')
        edges.append([i, j, dist])

# 가장 비용이 낮은 간선부터 비교해야하므로 비용 순으로 정렬해준다.
edges.sort(key=lambda x: x[2], reverse=True)

parent = list(range(N)) # 부모 노드 초기화, 초기값은 각 노드의 번호
edge_count = 0
cost = 0

# 최소가 되려면 간선이 (노드 수 - 1) 개만 있어야한다.
while edge_count < N-1:
    curr = edges.pop() # 가장 낮은 비용의 별과 비용 pop
    if find(curr[0]) != find(curr[1]): # 두 별이 아직 연결되지 않았다면
        union(curr[0], curr[1]) # 연결
        cost += curr[2] # 전체 비용에 추가
        edge_count += 1 # 간선 1회 추가

print(float(cost))