[백준 2887번] 파이썬 - 행성 터널

백준 2887_행성 터널

시간 제한 1초, 메모리 제한 128MB

# 조건

• 때는 2040년, 이민혁은 우주에 자신만의 왕국을 만들었다.
• 왕국은 N개의 행성으로 이루어져 있다.
• 민혁이는 이 행성을 효율적으로 지배하기 위해서 행성을 연결하는 터널을 만들려고 한다.
• 행성은 3차원 좌표위의 한 점으로 생각하면 된다.
  • 두 행성 A(xA, yA, zA)와 B(xB, yB, zB)를 터널로 연결할 때 드는 비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이다.
• 민혁이는 터널을 총 N-1개 건설해서 모든 행성이 서로 연결되게 하려고 한다.
• 이때, 모든 행성을 터널로 연결하는데 필요한 최소 비용을 구하는 프로그램을 작성하시오.

입력

• 첫째 줄에 행성의 개수 N이 주어진다. (1 ≤ N ≤ 100,000)
• 다음 N개 줄에는 각 행성의 x, y, z좌표가 주어진다.
• 좌표는 -109보다 크거나 같고, 109보다 작거나 같은 정수이다.
• 한 위치에 행성이 두 개 이상 있는 경우는 없다.

출력

• 첫째 줄에 모든 행성을 터널로 연결하는데 필요한 최소 비용을 출력한다.

# 접근 방법

MST 트리 - KRUSKAL 알고리즘 사용

• 간선을 만들기 위하여 한 행성에서 다른 행성으로 가는 각자의 최소비용을 저장해주는 반복문만 돌려도 시간초과가 발생한다 -> O( 100000 * 100000 )
• 따라서, 선택될 가능성이 있는 간선만 구성해주어야 한다.
• 각 행성마다의 최소거리는 -> 나와 가장 가까운 행성
  • 입력을 받을 때, X Y Z 좌표를 나눠서 저장해준다.
  • xl, yl, zl로 저장 후 sort 해준다.
• 거리별로 저장되어 있으므로, 해당 리스트들을 1부터 N까지 순회하며 비용을 구해준다.
  • 현재 인덱스와 직전 인덱스의 거리를 abs 값으로 구한 후 (abs, 직전 행성 번호, 현재 행성 번호) 로 저장해준 후 정렬해준다.
• 가중치가 가장 작은 간선부터 뽑아 KRUSKAL 알고리즘을 수행해준다.
  • 가중치가 가장 작은 간선부터 n-1개를 뽑아 mst 트리를 구성해준다.
  • 만약 사이클을 이룬다면, mst 트리를 구성할 수 없으므로, continue를 사용하고 다음 간선을 다시 선택해준다.

import sys
sys.stdin = open('input.txt')
input = sys.stdin.readline

def find(x):
    if x != parents[x]:
        parents[x] = find(parents[x])
    return parents[x]

def union(a, b):
    a = find(a)
    b = find(b)
    if a > b:
        parents[b] = a
    else:
        parents[a] = b

N = int(input())
cost = []
xl, yl, zl = [], [], []

# x, y, z의 차이 중 최소를 뽑는 것이기 때문에 행성 번호와 함께 따로 저장
for i in range(N):
    x, y, z = map(int, input().split())
    xl.append((x, i))
    yl.append((y, i))
    zl.append((z, i))

xl.sort()
yl.sort()
zl.sort()
for val in xl, yl, zl:
    for cur in range(1, N):
        now_planet = val[cur]
        pre_planet = val[cur - 1]
        # 비용, 현재 행성 번호, 연결시킨 행성 번호
        cost.append((abs(now_planet[0] - pre_planet[0]), pre_planet[1], now_planet[1]))
cost.sort()

parents = [i for i in range(N + 1)]
cnt, result = N - 1, 0
idx = -1
while cnt:
    idx += 1
    w, a, b = cost[idx]
    if find(a) == find(b):
        continue
    union(a, b)
    cnt -= 1
    result += w

print(result)