[BOJ 2887] 행성 터널 (Python)

문제

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

2887번: 행성 터널

 

풀이

크루스컬로 풀면 될 것 같은데 간선 정보가 없다. 

간선을 내가 만들어야 하는데 모든 경우의 수를 고려하면 시간초과 각이다 

 

비용이 각 양쪽 좌표 별 차이 값의 최소값이므로 좌표를 분리해서 간선을 만든다. 

x, y, z 각 좌표와 노드 번호를 저장 후 좌표 값 오름차순으로 정렬

좌표 별로 한쌍 씩 묶어서 비용과 노드를 edges 리스트에 저장 

edges 리스트 오름차순으로 정렬. 여기에는 (N-1)*3 개수의 간선이 저장되어 있다.

이제 크루스컬 알고리즘을 돌리면 된다. 

 

소스코드

import sys
input = sys.stdin.readline

n = int(input())
arr = [[0]*3 for _ in range(n)]

parent = [0]*n
for i in range(n):
    parent[i] = i

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

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

for i in range(n):
    data = list(map(int, input().split()))
    for j in range(3):
        arr[i][j] = data[j]

arr_x = []
arr_y = []
arr_z = []

for i in range(n):
    for j in range(3):
        if j == 0:
            arr_x.append((arr[i][j], i))
        if j == 1:
            arr_y.append((arr[i][j], i))
        if j == 2:
            arr_z.append((arr[i][j], i))

arr_x.sort()
arr_y.sort()
arr_z.sort()

edges = []
for i in range(1, n):
    edges.append((abs(arr_x[i][0] - arr_x[i - 1][0]), arr_x[i][1], arr_x[i-1][1]))
    edges.append((abs(arr_y[i][0] - arr_y[i - 1][0]), arr_y[i][1], arr_y[i-1][1]))
    edges.append((abs(arr_z[i][0] - arr_z[i - 1][0]), arr_z[i][1], arr_z[i-1][1]))

edges.sort()

minimum_cost = 0
for edge in edges:
    cost, a, b, = edge
    if find_parent(parent, a) != find_parent(parent, b):
        union(parent, a, b)
        minimum_cost += cost

print(minimum_cost)