[Python] 백준 / gold / 1774 : 우주신과의 교감

문제 링크 : https://www.acmicpc.net/problem/1774

문제 읽다 웃은건 처음이다 .. ㅋㅋ 하여튼 이건 크루스칼 알고리즘 문제. 우주신의 개수 N 은 1,000 이기 때문에 itertools 의 Combinations 로 가능한 모든 간선을 구해도 O(1,000,000) time 을 넘지 않는다.

그래서 다음과 같은 흐름으로 문제를 해결했다.
1. 이미 주어진 노드들을 Union 으로 연결
2. 가능한 edge 들을 combinations 로 모두 구해 거리 순으로 정렬
3. edge 를 탐색해나가면서 사이클 체크

기억할 것

1. 크루스칼 알고리즘은 Union-Find 를 활용해서 푼다.
2. 크루스칼 알고리즘은 사이클이 발생하는지 판단하면서 수행한다.
3. 사이클은 두 노드의 루트 노드가 같은지 체크함으로 판단할 수 있다.
4. round(num, 3) 은 num 을 소수점 셋째에서 반올림해 둘째자리까지의 수를 구해준다.
5. print("%.2f" num) 은 num 을 소수점 둘째자리까지 출력한다.

파이썬 코드

from itertools import combinations
import sys
import math

gods = dict()
already_connected = []

N, M = map(int, sys.stdin.readline().split())

for i in range(1, N+1):
    x, y = map(int, sys.stdin.readline().split())
    gods[i] = (x, y)

for _ in range(M):
    A, B = map(int, sys.stdin.readline().split())
    already_connected.append((A, B))

# 두 점을 이은 모든 조합
two_point_combinations = list(combinations(range(1, N+1), 2))

parent = [x for x in range(N+1)]


def find_parent(X):
    if parent[X] == X:
        return X
    else:
        parent[X] = find_parent(parent[X])
        return parent[X]


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

    if pa < pb:
        parent[pb] = pa
    else:
        parent[pa] = pb


def getDistance(S, T):
    x1, y1 = gods[S]
    x2, y2 = gods[T]
    return math.sqrt((x1-x2)**2 + (y1-y2)**2)


answer = 0

for A, B in already_connected:
    union(A, B)


edges = []
for A, B in two_point_combinations:
    # 노드, 노드, 거리
    edges.append((A, B, getDistance(A, B)))

edges.sort(key=lambda x: x[2])

for A, B, dist in edges:
    if find_parent(A) == find_parent(B):
        continue

    union(A, B)
    answer += getDistance(A, B)

answer = round(answer, 2)
print("%.2f" %answer)