파이썬으로 MST 알고리즘 구현하기

Prim 알고리즘

import heapq

n = 3
adj = [[],
       [(1,2), (3,3)],
       [(1,1), (2,3)],
       [(2, 2), (3,1)]
       ]

# Prim 알고리즘
visited = [False] * (n+1)
queue = list()
queue.append((0, 1))
answer = 0

while queue:
    cost, node = heapq.heappop(queue)
    if visited[node]:
        continue
    visited[node] = True
    answer += cost

    for next_cost, next_node in adj[node]:
        if not visited[next_node]:
            heapq.heappush(queue, (next_cost, next_node))

# MST 가중치 출력
print(answer)

실행 결과로 3이 출력된다.

Kruskal 알고리즘

# Kruskal 알고리즘
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

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

n = 3
parent = [i for i in range(n+1)]
edges = [(1, 1, 2), (2, 3, 2), (3, 1, 3)]

answer = 0
for c, a, b in edges:
    if union(a, b):
        answer += c

print(answer)

위와 마찬가지로 실행 결과로 3이 출력된다.