개요

• 크루스칼 알고리즘, 프림 알고리즘 2가지로 풀이 가능하다.

크루스칼 알고리즘

• 가중치가 가장 작은 간선부터 하나씩 이어가면서 모든 정점을 방문한다.
• 사이클이 발생 시, 그 간선은 건너뛴다.
  • 이를 판단하기 위해 유니온 파인드 구조를 사용한다.
    • 두 정점의 루트 노드가 동일한 경우 사이클이 생긴다고 판단한다.
    • 같지 않은 경우, 루트 노드를 합친다.
• 시간복잡도 O(ElogE)

# 크루스칼로 구현

v, e = map(int, input().split())

edges = [list(map(int, input().split())) for _ in range(e)]
edges.sort(key=lambda x: x[-1]) # 가중치 기준 정렬해주기

parent = [i for i in range(v+1)] # 초기화

# print(edges)
# print(parent)

def get_parent(x):
    if parent[x] == x: # 부모가 자기 자신이다 -> 루트이다.
        return x
    parent[x] = get_parent(parent[x])
    return parent[x]

def union_parent(a, b): # 수가 더 큰 쪽의 부모를, 작은쪽의 부모로 바꾼다 -> 두 간선을 합친다.
    a_root = get_parent(a)
    b_root = get_parent(b)
    if a_root < b_root:
        parent[b_root] = a_root
    else:
        parent[a_root] = b_root

def same_parent(a, b):
    return get_parent(a) == get_parent(b)

result = 0

for a, b, cost in edges:
    if not same_parent(a, b): # 부모가 같다 -> 사이클이 형성되는 것. 부모가 같지 않은 경우에만 한다.
        union_parent(a, b)
        result += cost

print(result)

프림 알고리즘

• 정점을 기준으로 인접한 간선들 중 가중치가 가장 작은 간선의 정점으로 이동한다.
  • 방문한 정점으론 이동하지 않는다.
• 우선순위 큐를 써서 구현하고, bfs와 비슷한 방식이다.
• 시간복잡도 O(ElogV)

import sys
import heapq
# 프림으로 구현
v, e = map(int, sys.stdin.readline().split())

edges = [[] for i in range (v+1)]
visited = [0] * (v+1)

def prim() :
    ret = 0
    min_heap =[]
    heapq.heappush(min_heap, (0, 1)) # 최소 스패닝 트리는 어디서 시작하든 상관없다. (애초에 전제가 모든 vertex를 다 돌아야 하는 거니깐.)
    
    while min_heap :
        cost, vertex = heapq.heappop(min_heap)
        
        if visited[vertex] : # 방문검사를 아래에서 하긴 했지만, 가중치 더 작은 edge에 의해 해당 vertex가 이미 선점되었을 수 있다. 즉 한번 더해준다.
            continue
        visited[vertex] = 1

        ret += cost
        
        for edge in edges[vertex] :
            if visited[edge[1]] == 1 :
                continue
            heapq.heappush(min_heap, edge) # 현재 만들어진 트리에서 갈 수 있는 경우의 수를 하나씩 더해나간다고 생각하자. 우리는 그중에서 가중치가 가장 작은 것들만 뽑을거다.
        
    return ret

for _ in range (e) :
    a, b, c = map(int, sys.stdin.readline().split())
    edges[a].append((c, b))
    edges[b].append((c, a))

print(prim())

참고

https://huiung.tistory.com/76

크루스칼 VS 프림 파헤치기