문제

[백준] 행성 터널 (Platinum 5)

요약

N개의 행성을 연결하는 터널을 만드려고 한다. 두 행성 A(xA, yA, zA)와 B(xB, yB, zB)를 터널로 연결할 때 드는 비용은 min( x(A)-x(B) , y(A)-y(B) , z(A)-z(B) )이다. 터널을 총 N-1개 건설해서 모든 행성이 서로 연결되게 하려고 할 때, 모든 행성을 연결하기 위한 최소 비용을 고르시오.

분류

  • 최소 스패닝 트리 (MST)
  • 정렬

풀이

1. 내 풀이 (MST)

원래 아래와 같은 방식으로 정렬을 진행해서 메모리 초과가 났다. N의 최댓값은 100,000 이었기 때문에 다른 방법을 고려해야 했다.

planets = [tuple(map(int, input().split())) for _ in range(n)]

edges = []
for i in range(n - 1):
    for j in range(i + 1, n):
        edges.append(
            (i, j, min(abs(planets[i][0] - planets[j][0]), abs(planets[i][1] - planets[j][1]), abs(planets[i][2] - planets[j][2]))))
min( x(A)-x(B) , y(A)-y(B) , z(A)-z(B) ) 이 부분의 연산에서, 두 좌표의 차가 작으려면 정렬된 상태에서 앞 뒤의 좌표만 고려해주면 된다는 발상을 했다. 그래서 x, y, z 별로 정렬을 하고, 순회를 통해 edge를 생성한다. 그리고 최소 스패닝 트리 알고리즘을 그대로 실행하면 된다. 
import sys
input = sys.stdin.readline

n = int(input())
planets = [[i] + list(map(int, input().split())) for i in range(n)]

edges = []
x = sorted(planets, key=lambda x: x[1])
for i in range(n - 1):
    edges.append((x[i][0], x[i + 1][0],
                 x[i + 1][1] - x[i][1]))
y = sorted(planets, key=lambda x: x[2])
for i in range(n - 1):
    edges.append((y[i][0], y[i + 1][0],
                 y[i + 1][2] - y[i][2]))
z = sorted(planets, key=lambda x: x[3])
for i in range(n - 1):
    edges.append((z[i][0], z[i + 1][0],
                 z[i + 1][3] - z[i][3]))

parent = [i for i in range(n)]


def find(idx):
    if parent[idx] != idx:
        parent[idx] = find(parent[idx])
    return parent[idx]


def union(s1, s2):
    if s1 < s2:
        parent[s2] = s1
    else:
        parent[s1] = s2


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

ans = 0

for v1, v2, dist in edges:
    s1 = find(v1)
    s2 = find(v2)
    if s1 != s2:
        ans += dist
        union(s1, s2)

print(ans)