[BOJ 2887] 행성터널 (Python)

최소 스패닝 트리 문제.

N이 100000이라서 모든 간선을 사용할 경우 메모리 초과가 발생한다.

간선을 최대한 줄이는 아이디어를 생각해내는 것이 핵심이었는데, 많이 헤맸다.

결국 x,y,z별로 한번씩 정렬을 해서 인접한 정점끼리만 계산해주면 된다는 결론이 나온다.

간선의 개수를 3 * (N-1)개 까지 줄일 수 있어서 메모리 초과를 피할 수 있다.

 

유니온파인드 + 크루스칼 알고리즘을 확실하게 구현할 줄 알아야 풀수 있다.

 

크루스칼 알고리즘의 핵심을 기억해놓자

간선 비용을 오름차순으로 정렬한 뒤 하나씩 선택해서

사이클이 발생하지 않을때만 최소 스패닝 트리에 추가한다.

 

import sys

input = sys.stdin.readline
n = int(input())
cost = []
loc = []
parent = [0]


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


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


for i in range(1, n + 1):
    x, y, z = map(int, input().split())
    loc.append((x, y, z, i))


def calc(i, j):
    return min(abs(i[0] - j[0]), abs(i[1] - j[1]), abs(i[2] - j[2]))


for i in range(3):
    sort_loc = sorted(loc, key=lambda x: (x[i]))
    for j in range(n - 1):
        cost.append((calc(sort_loc[j], sort_loc[j + 1]), sort_loc[j][3], sort_loc[j + 1][3]))

for i in range(1, n + 1):
    parent.append(i)

cost.sort()

result = 0
cnt = 0

for i in cost:
    c, start, end = i[0], i[1], i[2]
    if find(start) != find(end):
        union(start, end)
        result += c
        cnt += 1
        if cnt == n - 1:
            break

print(result)