[BOJ/Python] 1922: 네트워크 연결

gold 4 / 1922: 네트워크 연결

문제 흐름

최소 비용으로 모든 노드를 연결시키는 문제로 MST 문제라고 할 수 있다.
만약 노드의 수만 주어지고 어떤 두 노드든 연결시킬 수 있었다면,
잠정적으로 nC2 = (n * (n-1)) / 2 개의 간선이 생길 수 있으므로
Prim 알고리즘을 이용하는 것이 더 나았을 것이다.
하지만 이 문제에서는 간선의 수가 1 ≤ M ≤ 100,000로 제한되어 있으므로
Kruskal 알고리즘을 이용하였다.

꽤나 정석적인 Kruskal 알고리즘이라서 특별히 설명할 부분은 없을 것 같고,
만약 kruskal 함수에서 graph 자체에 c가 가장 앞으로 오는 튜플을 넣음으로써
번거롭게 정렬 key를 지정할 필요 없이 정렬할 수도 있다.

코드

# 네트워크 연결
# MST

import sys
input = sys.stdin.readline


### Kruskal
# union-find
def find(parent, x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]

    return x

def union(parent, x, y):
    x_root = find(parent, x)
    y_root = find(parent, y)
    
    if x_root == y_root:
        return False
    parent[y_root] = x_root
    return True

# kruskal
def kruskal(v, graph):
    parent = [i for i in range(v+1)]
    graph.sort(key=lambda x: x[2])
    result = 0

    for a, b, weight in graph:
        if union(parent, a, b):
            result += weight

    return result


import heapq
from collections import defaultdict
### Prim
def prim(graph):
    v = len(graph)
    visited = [0 for _ in range(v+1)]
    pq = []

    # initial setting
    heapq.heappush(pq, (0, 1))
    result = 0

    # repitition
    while pq:
        weight, x = heapq.heappop(pq)

        if visited[x]:
            continue

        visited[x] = 1
        result += weight

        for nx, weight in graph[x]:
            if not visited[nx]:
                heapq.heappush(pq, (weight, nx))
    return result


if __name__ == "__main__":
    n = int(input())
    m = int(input())
    graph_kruskal = []
    graph_prim = defaultdict(list)
    for _ in range(m):
        a, b, c = map(int, input().split())
        graph_kruskal.append((a, b, c))
        graph_prim[a].append((b, c))
        graph_prim[b].append((a, c))

    answer_kruskal = kruskal(n, graph_kruskal)
    print(answer_kruskal)
    # answer_prim = prim(graph_prim)
    # print(answer_prim)