BOJ #2887 행성 터널

읽기 전

• 불필요한 코드나 잘못 작성된 내용에 대한 지적은 언제나 환영합니다.
• 개인적으로 배운 점을 정리한 글입니다.

문제 링크

BOJ #2887 행성 터널

문제 풀이

최소신장트리로 가볍게 풀려다가 시간 초과에 메모리 초과까지 당했다. 기존 스패닝 트리 문제에서는 노드에 초점을 맞췄다면 이 문제는 총 비용만 계산하면 되므로 x, y, z에 대해 각각의 거리를 계산한다. 그리고 최소 비용이므로 굳이 다른 노드들끼리는 필요 없고 거리 순으로 정렬한 뒤 인접한 원소의 거리만 계산하면 된다.

• 유니온-파인드 알고리즘 함수를 정의해둔다.
• x, y, z를 별도의 세 배열에 저장하고 튜플 형식으로 좌표만 넣어 어디 행성의 것인지만 표기한다.
• 간선을 관리하는 배열을 선언한 뒤 정렬된 x, y, z 좌표 배열의 인접 원소끼리의 거리를 구해 해당 좌표들에 대한 정보를 담아 간선 배열에 넣는다.
• 간선 배열을 정렬하고 큐로 변환한다.
• 간선의 개수가 N - 1개가 될 때까지 유니온-파인드 알고리즘을 수행한다. 유효한 간선일 경우 간선의 개수에 1을 더해준다.

python 코드

import sys
from collections import deque
input = sys.stdin.readline

def solve():
    def upward(buf, idx):
        parent = nodes[idx]
        if parent < 0:
            return idx
        buf.append(idx)
        return upward(buf, parent)

    def find(idx):
        buf = []
        result = upward(buf, idx)
        for b in buf:
            nodes[b] = idx
        return result

    def union(x, y):
        x, y = find(x), find(y)
        if x == y:
            return False
        if nodes[x] < nodes[y]:
            nodes[y] = x
        elif nodes[x] > nodes[y]:
            nodes[x] = y
        else:
            nodes[x] -= 1
            nodes[y] = x
        return True

    N = int(input())
    x, y, z = [], [], []
    for i in range(N):
        a, b, c = map(int, input().split())
        x.append((a, i))
        y.append((b, i))
        z.append((c. i))
    x.sort()
    y.sort()
    z.sort()
    r = []
    for i in range(1, N):
        r.append(x[i][0] - x[i - 1][0], x[i - 1][1], x[i][1])
        r.append(y[i][0] - y[i - 1][0], y[i - 1][1], y[i][1])
        r.append(z[i][0] - z[i - 1][0], z[i - 1][1], z[i][1])
    r.sort()
    r = deque(r)
    cnt = 0
    sol = 0
    nodes = [-1 for _ in range(N)]
    while True:
        dist, i, j = r.popleft()
        if union(i, j) is True:
            cnt += 1
            sol += dist
            if cnt == N - 1:
                break
    return sol

print(solve())