[Python] BOJ 2887 - 행성 터널

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

2887번: 행성 터널

# 'N개의 행성을 터널로 연결하는데 필요한 최소비용' 보자마자 MST가 떠올랐다.

 

# 처음에 N개의 행성 사이의 터널 건설비용을 모두 구한 뒤, 크루스칼 알고리즘을 사용하려 했다.

   그런데, 모든 터널 건설비용을 구하기 위해선 N^2번의 연산이 필요하지만 N이 최대 10만이므로

   이렇게 하면 시간초과가 발생한다.

 

# A와 B 사이의 터널 건설비용은 min(|xA-xB|, |yA-yB|, |zA-zB|)이므로 x좌표순으로 정렬한 뒤 

   1번과 2번 사이의 거리, 2번과 3번 사이의 거리, ... , N-1번과 N번 사이의 거리를 구해서

   배열에 넣고 똑같이 y좌표, z좌표 순으로 정렬한 뒤 배열에 넣은 뒤

   크루스칼 알고리즘을 사용해서 풀었다.

 

   * 시간복잡도는 처음에 planet을 정렬하는데 O(NlogN)

                                 x좌표, y좌표, z좌표순으로 정렬하는데 O(NlogN) * 3

                                 tunnel배열에 원소가 N-1개씩 3번들어오고, 정렬하는데 O(3Nlog3N) = O(NlogN)

                                 MST실행하는데 O(ElogE) = O(NlogN)

     이므로 O(NlogN)으로 예상된다.

 

import sys
input = sys.stdin.readline
def find(c):
	if par[c] == c:
		return c
	else:
		par[c] = find(par[c])
	return par[c]

def union(a, b):
	a, b = find(a), find(b)
	par[max(a, b)] = min(a, b)

def check(a, b):
	return find(a) == find(b)

N = int(input())
par = [i for i in range(N)]
planet = []
for i in range(N):
	x, y, z = map(int,input().split())
	planet.append([x, y, z, i])
planet.sort()

tunnel = []
for i in range(N-1):
	tunnel.append((abs(planet[i][0] - planet[i+1][0]), planet[i][3], planet[i+1][3]))
planet.sort(key = lambda x: x[1])
for i in range(N-1):
	tunnel.append((abs(planet[i][1] - planet[i+1][1]), planet[i][3], planet[i+1][3]))
planet.sort(key = lambda x: x[2])
for i in range(N-1):
	tunnel.append((abs(planet[i][2] - planet[i+1][2]), planet[i][3], planet[i+1][3]))
tunnel.sort()

answer = 0
for t in tunnel:
    cost, a, b = t
    if not check(a, b):
        union(a, b)
        answer += cost
print(answer)