[BOJ] 2887 - 행성 터널 (Python)

 

 

아이디어

 

크루스칼 알고리즘으로 최소 스패닝 트리를 구현하되, 고려하지 않아도 될 터널들은 제외하여 시간을 단축한다.

 

 

풀이

 

좌표 x, y, z 각각을 기준으로 행성들을 정렬하여 인접한 두 행성과 두 행성을 연결하는 데 필요한 비용만을 고려한다.

 

import sys

input = sys.stdin.readline

def solution():

    def find(n):
        temp = n
        while n != group[n]:
            n = group[n]
        group[temp] = n
        return n

    def union(n1, n2):
        n1 = find(n1)
        n2 = find(n2)
        if n1 > n2:
            n1, n2 = n2, n1
        group[n2] = n1

    N = int(input())
    planets = []
    for i in range(N):
        x, y, z = map(int, input().split())
        planets.append((x, y, z, i))
    group = [i for i in range(N)]
    edges = []
    for j in range(3):
        planets.sort(key=lambda x: x[j])
        for i in range(N-1):
            d = abs(planets[i+1][j]-planets[i][j])
            edges.append((d, planets[i][3], planets[i+1][3]))
    edges.sort()
    res = cnt = 0
    for d, i, j in edges:
        if find(i) != find(j):
            union(i, j)
            res += d
            cnt += 1
            if cnt >= N-1:
                break
    print(res)

solution()