[코테 문제 풀이] 백준 2887번 - 행성 터널 (파이썬)

문제 설명

때는 2040년, 이민혁은 우주에 자신만의 왕국을 만들었다. 왕국은 N개의 행성으로 이루어져 있다. 민혁이는 이 행성을 효율적으로 지배하기 위해서 행성을 연결하는 터널을 만들려고 한다.

행성은 3차원 좌표위의 한 점으로 생각하면 된다. 두 행성 A($x_A$, $y_A$, $z_A$)와 B($x_B$, $y_B$, $z_B$)를 터널로 연결할 때 드는 비용은 min(|$x_A$-$x_B$|, |$y_A$-$y_B$|, |$z_A$-$z_B$|)이다.

민혁이는 터널을 총 N-1개 건설해서 모든 행성이 서로 연결되게 하려고 한다. 이때, 모든 행성을 터널로 연결하는데 필요한 최소 비용을 구하는 프로그램을 작성하시오.

• 입력 첫째 줄에 행성의 개수 N이 주어진다. (1 ≤ N ≤ 100,000) 다음 N개 줄에는 각 행성의 x, y, z좌표가 주어진다. 좌표는 -109보다 크거나 같고, 109보다 작거나 같은 정수이다. 한 위치에 행성이 두 개 이상 있는 경우는 없다.
  
• 출력 첫째 줄에 모든 행성을 터널로 연결하는데 필요한 최소 비용을 출력한다.
  

문제 풀이

모든 행성을 연결한다는 정보에서 크루스칼 알고리즘을 떠올렸다. 그러나 크루스칼 알고리즘은 간선을 가중치 별로 정렬해야 하는데, 문제에서는 간선에 대한 정보가 주어지지 않는다. 그렇다면 간선을 어떻게 만들어야 할까? 문제에서는 행성과 행성의 연결을 min(|$x_A$-$x_B$|, |$y_A$-$y_B$|, |$z_A$-$z_B$|) 로 정의했다. 때문에 x, y, z 축 각각을 기준으로 행성들을 정렬한 다음, 이웃하는 행성과 행성을 하나의 간선으로 취급한다. 각 축 별로 생성된 간선들을 가중치순으로 정렬한 뒤, 크루스칼 알고리즘을 수행하여 최소 신장 트리를 생성하여 문제를 풀 수 있다.

입출력 예시

예제 입력 1

5
11 -15 -15
14 -5 -15
-1 -1 -5
10 -4 -1
19 -4 19

예제 출력 1

4

CODE

import sys
input = sys.stdin.readline

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

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

def solution():
    N = int(input())
    nodes = []
    for i in range(N):
        X, Y, Z = map(int, input().split())
        nodes.append((i+1, X, Y, Z))
    
    # 1. x, y, z 축으로 정렬하여 이웃한 행성들의 edges 를 얻는다.
    edges = []
    # 1-1 dim축 기준으로 정렬 후 인접한 행성 edges에 추가
    for dim in range(1, 4): # 1:x, 2:y, 3:z
        sorted_list = sorted(nodes, key = lambda x: x[dim])
        for i in range(N-1):
            cost = abs(sorted_list[i][dim] - sorted_list[i+1][dim])
            v1 = sorted_list[i][0]
            v2 = sorted_list[i+1][0]
            edges.append((cost, v1, v2))
    
    # 2. 얻은 edges를 기반으로 크루스칼 알고리즘을 수행한다.
    parent = [i for i in range(N+1)]
    answer = 0
    for edge in sorted(edges):
        w, v1, v2 = edge
        
        # 2-1. v1, v2의 루트 노드가 다른 경우에만 Union 작업 실행
        if find_parent(parent, v1) != find_parent(parent, v2):
            union_parent(parent, v1, v2)
            answer += w
    
    return answer

if __name__ == "__main__":
    answer = solution()
    print(answer)

시간 복잡도