16398. 행성 연결 (Python)

16398번: 행성 연결

 

풀이

 크루스칼 알고리즘을 사용하여 최소 스패닝 트리를 만들면 된다.

 크루스칼 알고리즘에 대해서 간단히 설명하자면 입력으로 주어진 간선 정보를 비용을 기준으로 오름차순 정렬한다(우선수위 큐를 사용하면 된다). 비용이 낮은 순으로 간선을 꺼내어 해당 간선을 사용했을 때 그래프에 사이클이 발생하는지 체크한다. 사이클이 발생하지 않을때만 해당 간선을 사용하여 그래프를 완성해 나간다.

 이를 반복하여 총 N개의 정점이 있을 때 N-1개의 간선을 사용했다면 모든 정점이 사이클 없이 연결됨을 확인할 수 있다. 

 

코드

import sys
import heapq
input = sys.stdin.readline


# 크루스칼 알고리즘에 사용할 find, union 함수
def find(a):
    if p[a] == a:
        return a

    p[a] = find(p[a])
    return p[a]


def union(a, b):
    a = find(a)
    b = find(b)

    p[a] = b


if __name__ == '__main__':
    N = int(input())
    costs = []
    p = [i for i in range(N)]  # p[n] : n번 정점과 연결되어 있는 정점들의 꼭대기
    q = []

    for _ in range(N):
        costs.append(list(map(int, input().split())))

    for i in range(N):
        for j in range(i+1, N):  # 간선 정보들을 최소힙에 넣어준다.
            heapq.heappush(q, (costs[i][j], i, j))

    res, count = 0, 0
    while count != N - 1:  # N-1개의 간선을 사용했을 때 종료
        cost, a, b = heapq.heappop(q)

        if find(a) != find(b):
            union(a, b)
            res += cost
            count += 1

    print(res)