[백준 2887] 행성 터널 - Python

문제 보기

[사용한 알고리즘]

크루스칼(kruskal)

 

[문제 접근]

모든 행성은 연결되어 있기 때문에 x, y, z 좌표별로 정렬한 후 각 행성 사이의 거리를 구합니다. 거리를 모두 구한 후 가장 작은 거리부터 하나씩 연결하는 크루스칼 알고리즘을 사용하여 최소 스패닝 트리를 구하였습니다. 만약 정렬을 통해 각 행성 사이의 거리를 구하지 않고 모든 행성 사이의 거리를 구하면 시간 초과가 발생할 것이라 생각합니다.

 

[알고리즘]

1. 행성 좌표를 입력받습니다.

2. x, y, z 좌표별로 각각 정렬한 후 각 행성 사이의 거리를 구합니다.

3. 가장 작은 거리부터 사이클을 확인한 후 연결하는 크루스칼 알고리즘을 사용하여 최소 스패닝 트리를 구합니다.

 

[코드]

import sys


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


def union_parent(a, b, parent):
    a = find_parent(a, parent)
    b = find_parent(b, parent)
    if a < b:
        parent[b] = a
    else:
        parent[a] = b


if __name__ == "__main__":
    # input
    n = int(sys.stdin.readline())
    c = []
    for i in range(n):
        x, y, z = map(int, sys.stdin.readline().split())
        c.append((x, y, z, i))  # 좌표값, 인덱스

    # 각 좌표별로 정렬한 후 거리를 계산해서 edge를 만든다.
    edges = []
    for i in range(3):
        c.sort(key=lambda x:x[i])
        for j in range(1, n):
            edges.append((abs(c[j - 1][i] - c[j][i]), c[j - 1][3], c[j][3]))

    parent = [i for i in range(n)]
    total = 0
    edges.sort()

    for i in range(len(edges)):
        if find_parent(edges[i][1], parent) != find_parent(edges[i][2], parent):
            union_parent(edges[i][1], edges[i][2], parent)
            total += edges[i][0]
    print(total)