백준 2887번 : 행성 터널 in Python

2887번: 행성 터널

 

www.acmicpc.net

이 문제는 행성 간의 터널을 만들어 모든 행성들을 연결하고자 할 때 드는 비용의 최솟값을 구하는 문제이다. 이 때 각 행성의 위치는 3차원 좌표로 나타내고 두 행성 간 터널 건설 비용은 두 행성의 x좌표, y좌표, z좌표의 차이 중 가장 작은 값이다.

만약 행성의 개수가 적은 경우, 이전의 별자리 만들기 문제처럼 모든 행성 간의 터널 건설 비용을 구한 다음 가장 적은 비용의 터널부터 순차적으로 추가해 나가면 된다. 그러나 이 문제는 행성의 개수가 최대 10^5개까지 가능하기 때문에 이러한 방식으로 문제를 해결하려 하면 매우 오랜 시간이 소요된다. 그러므로 이보다 더 효율적으로 문제를 해결할 방법을 찾아야 한다.

백준 4386번 : 별자리 만들기 in Python

4386번: 별자리 만들기 www.acmicpc.net 이 문제는 2차원 평면 위의 별들을 이용해 별자리를 만들 때 선으로 연결된 별들의 거리의 합의 최소값을 구하는 문제이다. 즉, 이 문제는 최소 스패닝 트리를 구하는 문제..

wanna-be-developer-yjh.tistory.com

이 문제의 핵심은 이전의 문제와는 달리 터널의 건설 비용이 x좌표, y좌표, z좌표의 차이 중 가장 작은 값이라는 점이다. 이를 이용해 모든 행성 간의 터널 건설 비용을 따지지 않고, 가능한 터널들 중 일부분만 고려한다.

우선 각 좌표에 대해 모든 행성들을 정렬한다. 예를 들어 x좌표에 대해 모든 행성들을 정렬했을 때 정렬된 행성들의 x좌표들은 오름차순으로 나열된다. 이때 임의의 행성에 대해 그 행성과 정렬된 배열 내에서 이웃한 행성이 x좌표의 차이가 가장 작을 것이다. 그러므로 이웃한 행성은 해당 행성과 연결된 가능성이 있는 행성이고, 따라서 이 행성과의 터널을 heap에 추가한다. 이외의 나머지 터널은 x좌표의 차이가 더 크기 때문에 고려할 필요가 없다.(물론 y좌표 또는 z좌표의 차이가 작을 수도 있지만, 이는 y좌표와 z좌표에 대해서 정렬했을 때 고려하므로 넘어가도 된다.) 이를 y좌표와 z좌표에 대해서도 적용하면 총 고려할 터널의 개수는 3(N-1)개로 모든 터널의 개수인 N(N-1)개보다 훨씬 적은 값이 된다. 즉, 모든 터널을 고려하지 않고도 비용이 최소가 되는 트리를 찾을 수 있어 훨씬 효율적이다.

이를 코드로 구현하면 다음과 같다.

import sys
import heapq
input = sys.stdin.readline
sys.setrecursionlimit(10**5)

def find(u):
    if parent[u] == u: return u
    parent[u] = find(parent[u])
    return parent[u]
def union(u, v):
    root_u = find(u)
    root_v = find(v)
    parent[root_u] = root_v
# 고려할 터널을 heap에 추가하는 함수
def add_tunnels(dim):
    global N
    # 각 좌표에 대해 planets 배열을 정렬한다.
    planets_sorted = sorted(planets, key=lambda planet: planet[dim])
    for i in range(N-1):
        # planets_sorted 내의 이웃한 행성 간의 터널을 heap에 추가한다.
        u = planets_sorted[i][3]
        v = planets_sorted[i+1][3]
        dis = sys.maxsize
        for j in range(3):
            dis = min(dis, abs(planets_sorted[i+1][j] - planets_sorted[i][j]))
        heapq.heappush(heap, [dis, u, v])
N = int(input())
parent = [i for i in range(N+1)]
planets = []
heap = []
for i in range(1, N+1):
    x, y, z = map(int, input().split())
    planets.append([x, y, z, i])
for i in range(3):
    add_tunnels(i)
W = 0
while len(heap) > 0:
    dis, u, v = heapq.heappop(heap)
    if find(u) != find(v):
        union(u, v)
        W += dis
print(W)