BOJ - 2887 - 행성 터널

import sys


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

    parent[x] = find(parent, parent[x])
    return parent[x]


def unionx(parent, x, y):
    rootX = find(parent, x)
    rootY = find(parent, y)

    if not rootX == rootY: parent[rootY] = rootX


def solution():
    N = int(sys.stdin.readline())
    parent = [i for i in range(N)]
    location = []
    vertex = []

    for idx in range(N):
        x, y, z = map(int, sys.stdin.readline().split())
        location.append([x, y, z, idx])

    for j in range(3):
        location.sort(key=lambda x:x[j])
        tempLocation = location[0][3]
        for i in range(1, N):
            curLocation = location[i][3]
            vertex.append([
                abs(location[i][j] - location[i-1][j]),
                tempLocation, curLocation
            ])
            tempLocation = curLocation

    vertex.sort(key=lambda x: x[0])
    count = 0
    result = 0

    for weight, start, end in vertex:
        if find(parent, start) != find(parent, end):
            result += weight
            unionx(parent, start, end)

        if count == N - 1:
            break

    print(result)


solution()

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