[백준/BOJ] platinum5 - 2887번 행성 터널 (Python)

▶2887 - 행성 터널

▶문제

때는 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보다 작거나 같은 정수이다. 한 위치에 행성이 두 개 이상 있는 경우는 없다. 

 

▶출력

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

 

▶풀이

크루스칼을 이용해서 문제를 풀어보았다.

union, find를 사용해서 root가 가리키는 것이 무엇인지 확인해주면 cycle 없이 풀 수 있을 것이다.

 

그렇게 풀었더니... 메모리 초과

import heapq


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


n = int(input())
li = []
dist_li = {}
root = list(range(n))
edge, cost = 0, 0

for i in range(n):
    li.append(list(map(int, input().split())))
for i in range(n):
    for j in range(i + 1, n):
        dist_li[i, j] = min(abs(li[i][0] - li[j][0]),
                            abs(li[i][1] - li[j][1]),
                            abs(li[i][2] - li[j][2]))
dist_li = sorted(dist_li.items(), key=lambda item: item[1])

for point, weight in dist_li:
    a, b = find(point[0]), find(point[1])
    if a != b:
        root[b] = a
        edge += 1
        cost += weight
    if edge == n-1:
        break

print(cost)

간선을 n(n - 1) / 2만큼 다 확인하고 만들어준다.

그래서 너무 큰 메모리를 차지하게 되었고, 그게 아니어도 시간 초과가 날 것 같았다.

 

곰곰이 생각해보았는데

i, j행성 사이에 거리는 x, y, z좌표의 차이중 작은 값을 알면 된다.

그럼 x, y, z별로 반복문을 돌리고, 그 차이를 각각 저장해주면

3 * (n - 1)이 될 거 같았다.

정렬하면 어차피 작은 값이 앞으로 가게 되고, 뒤에서 중복 선택되지 않을 것이기 때문이다.

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


n = int(input())

li = []
dist_li = []
root = list(range(n))
edge, cost = 0, 0

for i in range(n):
    x, y, z = list(map(int, input().split()))
    li.append([x, y, z, i])

for j in range(3):
    li = sorted(li, key=lambda x: x[j])
    bf_location = li[0][3]
    for i in range(1, n):  # 각 좌표별로 간선추가
        cur_location = li[i][3]
        dist_li.append([abs(li[i][j] - li[i - 1][j]), bf_location, cur_location])
        bf_location = cur_location

dist_li = sorted(dist_li, key=lambda x: x[0])
for point, s, f in dist_li:
    a, b = find(s), find(f)
    if a != b:
        root[b] = a
        edge += 1
        cost += point
    if edge == n - 1:
        break

print(cost)

흠........ 일단 pypy3로 제출하니 정답은 나왔다...

하지만 찝찝한 점은 python3으로 제출했을 때는 시간 초과가 났다는 점이다ㅠ

이 부분에 대해서는 조금 더 공부를 해서 다시 제출해보아야겠다.

 

pypy3가 python3보다 시간적으로는 좋다는 걸 알고 있어서 일단 pypy3로 제출하였다.