[백준] 2887번 행성 터널 (by Python)

문제

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

2887번: 행성 터널

 

풀이

문제를 보고 사이클을 생성하지 않는 범위 내에서 비용이 적은 순으로 edge를 (v - 1)개까지 선택하는 크루스칼 알고리즘이라는 걸 알 수 있다. 나는 처음에 min(|xA-xB|, |yA-yB|, |zA-zB|) 라고 제시된 각 edge의 비용을 모든 A와 B에 대해 계산하여 시간 초과를 마주하게 되었다. 그리고 좀 더 고민해본 결과 모든 좌표들의 x값, y값, z값을 각각 별도로 정렬한 후 각 축에서 바로 옆에 있는 좌표와의 거리만 edge들만 고려하면 된다는 사실을 알 수 있었다. 간단한 예로, x좌표가 각각 (1, 3, 8)인 점 3개가 주어졌을 때 비용 7을 들여 첫 번째 점과 세 번째 점을 연결하는 것보다 같은 비용으로 첫 번째 점과 두 번째 점을, 두 번째 점과 세 번째 점을 잇는 것이 더 효율적이기 때문이다. 해당 아이디어를 바탕으로 Union-Find 방식을 사용해 크루스칼 알고리즘을 구현했다. 코드는 다음과 같다.

 

import sys
input = sys.stdin.readline

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

def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b

n = int(input())
parent = list(range(n))
x = []
y = []
z = []

for i in range(n):
    a, b, c = map(int, input().split())
    x.append((a, i))
    y.append((b, i))
    z.append((c, i))
x.sort()
y.sort()
z.sort()
edges = []
for i in range(n - 1):
    edges.append((x[i + 1][0] - x[i][0], x[i][1], x[i + 1][1]))
    edges.append((y[i + 1][0] - y[i][0], y[i][1], y[i + 1][1]))
    edges.append((z[i + 1][0] - z[i][0], z[i][1], z[i + 1][1]))

edges.sort()
answer = 0

for e in edges:
    cost, a, b = e
    if find_parent(a) != find_parent(b):
        union_parent(a, b)
        answer += cost
print(answer)