백준 2887 행성 터널

문제

행성 터널

행성 터널을 연결하는데 드는 최소 비용을 구하자

난이도 : Platinum5

---

풀이

1. 각 행성 좌표의 x, y, z축을 나눠 3개의 리스트에 각각 저장한다.
2. 이 때 행성 번호 저장을 위해 [행성 번호, 좌표] 형태로 저장하며 2차원 배열 형태가 된다.
3. 각 축 리스트를 좌표 기준으로 정렬한다. (인접한 행성 쌍만 edges에 포함시키기 위해)
4. 정렬 완료된 리스트에서 i와 i+1번째 행성 번호, 두 행성 간 cost를 하나의 튜플로 묶어 edges에 포함시킨다.
5. x, y, z축에 대해 위 연산을 수행하면 edges에 총 3n개의 데이터가 들어가게 된다.

---

코드

import sys

n = int(sys.stdin.readline())
data = [list(map(int, sys.stdin.readline().split())) for _ in range(n)]
# 행성 각 축을 따로 저장하고, 행성 번호와 행성 좌표를 묶어 2차원 배열로 저장
planet_x = list(map(lambda x: x[0], data))
planet_y = list(map(lambda x: x[1], data))
planet_z = list(map(lambda x: x[2], data))
planet_x = list(map(lambda x: [x[0], x[1]], enumerate(planet_x)))
planet_x.sort(key=lambda x: x[1])
planet_y = list(map(lambda x: [x[0], x[1]], enumerate(planet_y)))
planet_y.sort(key=lambda x: x[1])
planet_z = list(map(lambda x: [x[0], x[1]], enumerate(planet_z)))
planet_z.sort(key=lambda x: x[1])

edges = []
# 인접한 행성 쌍끼리 edges에 포함시키고 kruskal을 수행하면 최소 cost를 구할 수 있다.
for i in range(len(planet_x) - 1):
    edges.append((planet_x[i][0], planet_x[i+1][0], abs(planet_x[i][1] - planet_x[i+1][1])))
for i in range(len(planet_y) - 1):
    edges.append((planet_y[i][0], planet_y[i+1][0], abs(planet_y[i][1] - planet_y[i+1][1])))
for i in range(len(planet_z) - 1):
    edges.append((planet_z[i][0], planet_z[i+1][0], abs(planet_z[i][1] - planet_z[i+1][1])))

parent = [i for i in range(n)]


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


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


def kruskal():
    total_cost = 0
    edges.sort(key=lambda x: x[-1])

    for i in range(len(edges)):
        a, b, cost = edges[i]
        if find(a, parent) != find(b, parent):
            union(a, b, parent)
            total_cost += cost
    print(total_cost)

kruskal()

        

---

결과

처음엔 인접 행성 쌍이 아닌 모든 행성 쌍에 대해서 edges를 만들자 메모리 초과가 났다.
축을 따로 뽑고 인접 행성 쌍으로 edges를 생성한다는 생각을 떠올리는 것이 제일 어려운 문제였던 것 같다.