백준 2887 <행성 터널> Python

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

---

최소 스패닝 트리의 응용 문제입니다.

해당 문제의 가장 큰 문제는 최소 스패닝 트리를 찾는 알고리즘 자체보다는 어떤 엣지를 타겟으로 삼아야할지 정하는 것 입니다. N==100,000 이라고 한다면 모든 엣지를 탐색하기에는 시간이 부족하거나 메모리가 부족할 것이기 때문입니다.

그렇다면 어떤 엣지가 최소 스패닝 트리를 구성하는 엣지가 될 수 있는지 생각해보아야 하는데, 단순히 생각해보면 가장 인접해있는 점이 그 후보에 가장 가깝다는 것을 알 수 있습니다. 가장 거리가 가까울테고, 가장 가까운 점은 하나일테니 어떤 두 점이 연결되었다면 다른 점에 연결될 가능성이 없어 순환이 발생하지도 않습니다.

다만 해당 문제에서 엣지의 코스트는 두 점의 절대적인 거리가 아닌 x y z 의 좌표 차 중 최솟값이기 때문에 x y z 좌표에 각각 가장 가까운 점들을 후보 엣지로 등록할 수 있습니다.

---

import sys

sys.setrecursionlimit(10000)

input = sys.stdin.readline

## Union-Find

# Union

def Union(node1, node2, group_list):

    # 각 노드의 그룹

    group1 = Find(node1, group_list)

    group2 = Find(node2, group_list)

    # 두 그룹이 같으면

    if group1 == group2:

        # 병합하지 않음

        return False, group_list

   

    # 두 그룹이 다르면

    # 작은 쪽으로 병합

    if group1 > group2:

        group_list[group1] = group2

   

    else:

        group_list[group2] = group1

   

    return True, group_list

# Find

def Find(node, group_list):

    # 그룹의 대표가 자신이 아니면

    if group_list[node] != node:

        # 재귀적으로 업데이트

        group_list[node] = Find(group_list[node], group_list)

   

    return group_list[node]

def solution(N, idxs):

    # 엣지 정리

    edges = []

    # x y z 로 각 정렬해 인접한 점들만 탐색

    for key in range(3):

        # 정렬된 좌표

        s_idxs = sorted(idxs, key=lambda x: x[key])

        for i in range(N-1):

            edges.append([s_idxs[i][-1], s_idxs[i+1][-1], min([abs(a-b) for a, b in zip(s_idxs[i][:3], s_idxs[i+1][:3])])])

    edges = sorted(edges, key=lambda x: x[-1])

    # 그룹 리스트

    group_list = list(range(N))

    # 코스트

    cost = 0

    # 크루스칼

    for a, b, c in edges:

        # 유니온

        state, group_list = Union(a, b, group_list)

        # 연결되었으면 코스트 더해주기

        if state:

            cost += c

   

    print(cost)

# 입력

N = int(input().strip())

idxs = [list(map(int, input().strip().split())) + [i] for i in range(N)]

solution(N, idxs)