Baekjoon 2887. 행성 터널 / Python

https://www.acmicpc.net/problem/2887

문제 설명

2040년, 이민혁은 우주에 자신만의 왕국을 만들었습니다. 왕국은 N개의 행성으로 이루어져 있으며, 이 행성들을 터널로 연결하여 효율적으로 지배하려고 합니다. 각 행성은 3차원 좌표계의 한 점으로 생각할 수 있고, 두 행성 A와 B를 연결하는 터널의 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)로 계산됩니다. 민혁이는 모든 행성을 최소 비용으로 연결하려고 합니다.

이 문제는 최소 신장 트리(MST)를 구하는 문제로, 크루스칼 알고리즘을 활용하여 해결할 수 있습니다.

알고리즘 설명

1. Union-Find (Disjoint Set):
   
   • union_find(x): x가 속한 집합의 대표자, 부모노드를 찾는다.
   • union_set(x, y): x와 y가 속한 집합을 병합
2. 입력 처리:
   
   • 각 행성의 좌표에 따라 정렬하여, 인접한 행성 간의 비용을 계산
3. 최소 힙 (Min-Heap):
   • 모든 간선 정보를 최소 힙에 저장, 이는 가장 작은 비용의 간선을 효율적으로 꺼내기 위함이다.
4. 크루스칼 알고리즘:
   • 가장 작은 비용의 간선을 선택하여 MST를 구성
   • 사이클을 방지하기 위해 union-find를 사용
   • N-1개의 간선이 선택될 때까지 반복
5. 출력:
   • MST를 구성하는데 필요한 최소 비용 출력

from heapq import heappop, heappush
import sys
input = sys.stdin.readline

# 재귀 한도 설정
sys.setrecursionlimit(10**8)

# Union-Find 함수 
def union_find(x):
    if x == parent[x]:
        return x
    parent[x] = union_find(parent[x])
    return parent[x]

def union_set(x, y):
    x, y = union_find(x), union_find(y)
    if x > y:
        x, y = y, x
    parent[y] = x

# 행성의 개수 
n = int(input())

# 부모 노드 초기화
parent = [i for i in range(n)]

# 행성의 좌표
arr = []
for i in range(n):
    x, y, z = map(int, input().split())
    arr.append((x, y, z, i))

# 좌표별 정렬
x_dist = sorted(arr, key=lambda x: x[0])
y_dist = sorted(arr, key=lambda x: x[1])
z_dist = sorted(arr, key=lambda x: x[2])

# 최소 힙
q = []

# 인접 행성 간의 간선 추가
for i in range(n - 1):
    heappush(q, (abs(x_dist[i][0] - x_dist[i + 1][0]), x_dist[i][3], x_dist[i + 1][3]))
    heappush(q, (abs(y_dist[i][1] - y_dist[i + 1][1]), y_dist[i][3], y_dist[i + 1][3]))
    heappush(q, (abs(z_dist[i][2] - z_dist[i + 1][2]), z_dist[i][3], z_dist[i + 1][3]))

res = 0

# 최소 신장 트리
def sol(cnt=0):
    global res
    if cnt == n - 1:
        print(res)
        exit(0)
    dist, a, b = heappop(q)
    if union_find(a) != union_find(b):
        union_set(a, b)
        res += dist
        sol(cnt + 1)
    sol(cnt)

sol()