[백준] 2887번 - 행성 터널 (파이썬)

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

# 행성 터널
import copy
import heapq
import sys
input = sys.stdin.readline

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

planets_x = copy.deepcopy(planets)
planets_y = copy.deepcopy(planets)
planets_z = copy.deepcopy(planets)
planets_x.sort(key=lambda index: index[0])
planets_y.sort(key=lambda index: index[1])
planets_z.sort(key=lambda index: index[2])

arr = [[] for _ in range(N)]
for i in range(1, N):
    a, b, c, i1 = planets_x[i-1]
    x, y, z, i2 = planets_x[i]
    cost = min(abs(a-x), abs(b-y), abs(c-z))
    arr[i1].append((cost, i2))
    arr[i2].append((cost, i1))

    a, b, c, i1 = planets_y[i-1]
    x, y, z, i2 = planets_y[i]
    cost = min(abs(a-x), abs(b-y), abs(c-z))
    arr[i1].append((cost, i2))
    arr[i2].append((cost, i1))

    a, b, c, i1 = planets_z[i-1]
    x, y, z, i2 = planets_z[i]
    cost = min(abs(a-x), abs(b-y), abs(c-z))
    arr[i1].append((cost, i2))
    arr[i2].append((cost, i1))


def prim(start):
    global result
    q = []
    heapq.heappush(q, (0, start))
    visited = [False] * N
    count = 0
    while q:
        if count == N:
            break
        cost, node = heapq.heappop(q)
        if not visited[node]:
            count += 1
            result += cost
            visited[node] = True
            for ncost, nnode in arr[node]:
                heapq.heappush(q, (ncost, nnode))


result = 0
prim(0)
print(result)

 

MST, 프림, 정렬

1. 행성의 정보를 planets리스트에 입력을 받는다

2. 정렬

    - palnets_x, planets_y, planets_z에 각각 planets를 deepcopy를 해준다 (원래 행성의 인덱스값 i도 같이)

    - 위 리스트를 각각 x좌표, y좌표, z좌표를 기준으로 정렬해준다.

4. 인접리스트에 각 좌표를 기준으로 가장 가까운 행성만을 append해준다.

5. 프림알고리즘 실행

    - 현재 방문한 노드와 연결된 간선중 가장 비용이 낮은 간선을 하나씩 포함시킨다.

6. result 출력

 

데이터 수가 많아서 정렬을 거치지 않으면 시간초과가 발생한다.

다른 행성과 연결하는데 발생하는 비용은 각 좌표차의 최솟값이기 때문에 좌표별로 정렬을 해서

가장 가까운 행성만을 인접리스트에 넣어주면 된다.