백준 1922번 : 네트워크 연결 in Python

1922번: 네트워크 연결

 

www.acmicpc.net

이 문제는 컴퓨터와 컴퓨터를 모두 연결하는 네트워크를 구축할 때 드는 비용의 최솟값을 구하는 문제이다.

이는 각 컴퓨터를 노드로, 컴퓨터 사이를 연결하는 선은 간선으로 볼 때 그 결과 나오는 그래프의 minimum spanning tree를 구하는 문제로 볼 수 있다. 따라서 MST를 구하기 위한 Kruskal's algorithm이나 Prim's algorithm을 이용하면 되는데, 여기서는 Kruskal's algorithm을 사용한다.

즉, 모든 노드마다 하나의 집합으로 간주한 뒤, 주어진 간선들을 비용이 적은 순서대로 저장한다. 그리고 간선을 하나씩 꺼내면서 간선의 양 끝 점이 같은 집합 안에 있지 않은 경우 연결하는 것을 반복한다. 이때 같은 집합에 있는지는 union find 함수를 이용한다.

이를 코드로 구현하면 다음과 같다.

import heapq
N = int(input())
M = int(input())
edges = []
for _ in range(M):
    a, b, c = map(int, input().split())
    heapq.heappush(edges, [c, a, b])
def find(u):
    if parent[u] != u: parent[u] = find(parent[u])
    return parent[u]
def union(u, v):
    root_u = find(u)
    root_v = find(v)
    if root_u != root_v: parent[root_u] = root_v
parent = [i for i in range(N+1)]
cost = 0
while edges:
    c, a, b = heapq.heappop(edges)
    if find(a) != find(b):
        union(a, b)
        cost += c
print(cost)