우주신과의 교감

1. 개요

• https://www.acmicpc.net/problem/1774

2. 크루스칼 알고리즘 풀이

from collections import deque
from sys import stdin
import math

def solution():
    q = deque()
    INPUT = lambda: stdin.readline().rstrip()
    READINTS = lambda: list(map(int, INPUT().split()))
    
    N, M = map(int, INPUT().split())
    GODS = [READINTS() for _ in range(N)]
    EDGES = set(tuple(n-1 for n in READINTS()) for _ in range(M))

    P = [n for n in range(N)]

    def calcEdge(a, b):
        ar, ac = GODS[a]
        br, bc = GODS[b]
        return (math.sqrt((ar-br)**2+(ac-bc)**2), a, b)
    
    def parent(n):
        if P[n]!=n:
            P[n] = parent(P[n])
        return P[n]

    def union(a, b):
        pa, pb = parent(a), parent(b)

        if pa==pb:
            return False
        elif pa<pb:
            P[pb] = pa
        else:
            P[pa] = pb

        return True

    def calcEdges():
        candidates = []

        for i in range(N):
            for j in range(i+1, N):
                if i==j:
                    continue

                if (i, j) in EDGES or (j, i) in EDGES:
                    candidates.append((0, i, j))
                else:
                    candidates.append(calcEdge(i,j))

        candidates.sort()
        return candidates

    def kruskal():
        traced = set()
        distances = 0

        for d, x, y in calcEdges():
            if not union(x, y):
                continue

            distances += d
            traced.add(x)
            traced.add(y)

            if len(traced)==N:
                return "{:.2f}".format(distances)
            
    print(kruskal())

solution()