[TIL] 최소 신장 트리(MST) 2 - 프림(prim) 알고리즘

프림 알고리즘(Prim's algorithm)

• 시작 정점을 선택한 후 정점에 인접한 간선 중 최소 가중치의 간선이 연결된 정점을 선택한다.
• 선택한 정점에서 위의 과정을 반복하며 최소 신장 트리를 확장해 나간다.

 

프림 알고리즘 구현

•  노드의 연결을 확인하기 위한 집합과 연결된 간선들을 관리하기 위한 간선 리스트를 사용한다.
• 임의의 정점을 선택 후 노드 집합에 삽입한다.
• 해당 정점에 연결된 간선들을 간선 리스트에 삽입한다.
• 간선 리스트에서 최소 가중치를 가지는 간선부터 추출한다. -> heapq 사용
• 만약, 해당 간선에 연결된 정점이 위의 노드 집합에 들어 있다면 다른 간선을 찾는다. (사이클 x)
• 그렇지 않다면 그 간선을 선택하고 mst를 리스트에 삽입한다.
• 추출된 간선은 간선 리스트에서 제거한다.
• 간선 리스트에 간선이 없을 때까지 계속 반복한다.

1. 간선 리스트 구현

edges=[
    (3, 'A', 'B'), (5, 'A', 'D'), (7, 'A', 'C'),
    (6, 'B', 'C'), (6, 'B', 'E'),
    (3, 'C', 'D'), (6, 'C', 'E'),
    (9, 'D', 'E')
]

 

2. 프림 main 함수 구현

from collections import defaultdict
from heapq import *

def prim(start_node, edges):
    mst = list()
    adj = defaultdict(list)		# value가 지정되어있지 않을 시 빈 리스트로 초기화
    for w, n1, n2 in edges:
        adj[n1].append((w,n1,n2))
        adj[n2].append((w,n2,n1))
        
    connected_node = set(start_node)	# 시작 노드를 연결 노드 집합에 삽입
    next_edge_list = adj[start_node]
    heapify(next_edge_list)				# 간선 리스트에서 최소 가중치를 가진 간선을 꺼내기 위해
    
    while next_edge_list:
        w, n1, n2 = heappop(next_edge_list)
        if n2 not in connected_node:
            connected_node.add(n2)
            mst.append((w, n1, n2))
            
            for edge in adj[n2]:			# n2 정점의 인접 간선 리스트 중 연결된 정점이 노드 집합에 없다면 heappush
                if edge[2] not in connected_node:
                    heappush(next_edge_list, edge)

    return mst

 

3. 실행 결과

prim('A', edges)

#[(3, 'A', 'B'), (5, 'A', 'D'), (3, 'D', 'C'), (6, 'B', 'E')]

 

프림 알고리즘의 시간복잡도?
- O(ElogE)