[알고리즘, 파이썬] 프림 알고리즘 - 2

이제 본격적으로 프림 알고리즘을 파이썬 코드로 작성해 보도록 하겠습니다! 

파이썬으로 프림 알고리즘을 간단하게 구현하기 위해서 두 가지 라이브러리를 소개 해드리겠습니다. 

 

참고) heapq 라이브러리 활용을 통해 우선순위 큐 사용하기

우선순위 큐는 우선 순위가 가장 높은 자료(data)를 가장 먼저 꺼낼 수 있는 자료 구조인데요, 파이썬에서는 이를 heapq라는 라이브러리로 정의해 놓았습니다! 이 자료구조는 삽입연산이 O(logN)으로 빠른 연산속도를 보여줍니다 

(자세한 내용은 여기서는 스킵) 

 

아무튼 요지는 이 라이브러리를 사용하면 가중치가 작은 것부터 우선적으로 쉽게 값을 참조하고 뽑아 쓸수 있어 프림알고리즘을 구현하는데 굉장이 유용하다는 것입니다! 

 

heapq에 내장된 함수인 heapify함수는 입력으로 받은 리스트 데이터를 자동으로 heap형태로 한번에 변환해줍니다! 

import heapq 

graph_data = [[2, 'A'], [5, 'B'], [3, 'C']]

 
heapq.heapify(graph_data)

for index in range(len(graph_data)):

    print (heapq.heappop(graph_data))

print (graph_data)

결과: [2, 'A'] [3, 'C'] [5, 'B'] []

참고) collections 라이브러리의 defaultdict 함수 활용하기

defaultdict 함수는 key에 대한 value를 지정하지 않았더라도 자동으로 빈 리스트로 초기화해주는 함수입니다. 

from collections import defaultdict

list_dict = defaultdict(list)
print (list_dict['key1'])

결과: []

 

프림 알고리즘 구현 

앞서 포스트에서 프림알고리즘 동작 방식 5단계를 기억하시나요? 이를 좀더 구현에 집중하여 살펴보겠습니다. 

 

1. 모든 간선 정보를 저장 (adjacent_edges)
2. 임의의 정점을 선택, '연결된 노드 집합(connected_nodes)'에 삽입
3. 선택된 정점에 연결된 간선들을 간선 리스트(candidate_edge_list)에 삽입
4. 간선 리스트(candidate_edge_list)에서 최소 가중치를 가지는 간선부터 추출해서,
   • 해당 간선에 연결된 인접 정점이 '연결된 노드 집합'에 이미 들어 있다면, 스킵함(cycle 발생을 막기 위함)
   • 해당 간선에 연결된 인접 정점이 '연결된 노드 집합'에 들어 있지 않으면, 해당 간선을 선택하고, 해당 간선 정보를 '최소 신장 트리(mst)'에 삽입
     • 해당 간선에 연결된 인접 정점의 간선들 중, '연결된 노드 집합(connected_nodes)' 에 없는 노드와 연결된 간선들만 간선 리스트(candidate_edge_list) 에 삽입
       • '연결된 노드 집합(connected_nodes)' 에 있는 노드와 연결된 간선들을 간선 리스트에 삽입해도, 해당 간선은 스킵될 것이기 때문임
       • 어차피 스킵될 간선을 간선 리스트(candidate_edge_list)에 넣지 않으므로 해서, 간선 리스트(candidate_edge_list)에서 최소 가중치를 가지는 간선부터 추출하기 위한 자료구조 유지를 위한 effort를 줄일 수 있음 (예, 최소힙 구조 사용)
5. 선택된 간선은 간선 리스트에서 제거
6. 간선 리스트에 더 이상의 간선이 없을 때까지 3-4번을 반복

 

프림 알고리즘 파이썬 코드

이제 위에서 작성한 글을 토대로 파이썬 코드를 구현해 보겠습니다

먼저 이전 포스트에서 예시로 든 그래프를 다음과 같이 리스트 형태로 저장하였다 가정합시다. 

#중복되는 간선들을 제외하고 표시함 
myedges = [
    (7, 'A', 'B'), (5, 'A', 'D'),
    (8, 'B', 'C'), (9, 'B', 'D'), (7, 'B', 'E'),
    (5, 'C', 'E'),
    (7, 'D', 'E'), (6, 'D', 'F'),
    (8, 'E', 'F'), (9, 'E', 'G'),
    (11, 'F', 'G')
]

 

이제 부터 위 그래프와 시작할 임의의 노드를 입력으로 받아서 최소신장트리 값을 반환하는 prim 이라는 함수를 작성해 보도록 하겠습니다! 

 

from collections import defaultdict 
from heapq import *

def prim(start_node, edges):
	mst = list() 
    adjacent_edges = defaultdict(list)
    for weight, n1, n2 in edges: 
    	adjacent_edges[n1].append((weight, n1, n2)) 
        adjacent_edges[n2].append((weight, n2, n1)) 
      
    connected_nodes = set(start_node) 
    candidate_edge_list = adjacent_edges[start_node] 
    heapify(candidate_edge_list)     

먼저 기본적인 변수 설정과 초기화 과정입니다.  최종적인 결과값(mst)를 반환할 리스트 mst, 

모든 쌍방 간선들의 정보를 저장할 adjacent_edges, 연결된 노드들을 저장할 connected_nodes,

후보군 간선들을 저장할 최소우선순위 큐로 정의된 candidate_edge_list  를 정의하여 줍니다

 

중간에 for 문으로 시작하는 부분이 이해가 안가실 수 도 있는데, 이는 어떤 노드를 시작으로 접근하더라도 해당하는 간선 정보를 불러올 수 있도록 모든 쌍방으로 연결된 간선을 저장하는 것입니다. 

 

입력으로 받는 myedges 그래프 같은 경우는 중복되는 간선이 없습니다. 

즉 (7, `A' , 'B') 라는 간선이 저장되어 있다면 (7, `B' , 'A') 라는 간선은 저장되어 있지 않다는 말입니다. 

 

하지만 for 문 연산을 통해 adjacent_edges에 그 쌍방의 간선정보까지 모두 저장해 줍니다.

왜 이런 짓을 하냐구요? 그래야지 어떤 노드의 값으로 접근하든 그 노드에 인접한 모든 간선의 값들을 조회할 수 있기 떄문 입니다.  만약에 그저 입력으로 들어온 그래프 값으로만 조회한다면 노드 B로 참조를 시도했을 때 

 (7, `B' , 'A') 라는 간선의 정보는 저장되어 있지 않기 때문에 인접 간선정보가 누락이 되버릴 겁니다. 

 

 

이제 초기화까지 완료하였으니 핵심 알고리즘을 구현해 보도록 하겠습니다.

while candidate_edge_list:
	weight, n1, n2 = heappop(candidate_edge_list)
    if n2 not in connected_nodes: 
    	connected_nodes.add(n2)
        mst.append((weight, n1, n2)) 
        
        for edge in adjacent_edges[n2]:
        	if edge[2] not in connected_nodes:
            	heappush(candidate_edge_list, edge) 

return mst
 

 

글로는 엄청 복잡했던 것 같은데, 구현하니 채 10줄 정도 밖에 나오지 않습니다;; 

구현은 위에서 작성한 구현 알고리즘을 그대로 따라 가시면 됩니다

 

먼저 candidate_edge_list 에 저장된 mst에 저장될 수 있는 후보 간선을 최소값이 작은것 부터 차례대로 불러옵니다. 

만약 (weight, n1->n2) 인 간선을 뽑았을 때 새롭게 잇게 되는 노드인 n2가 아직 connected_nodes에 저장되지 않았다면?  

 

해당 간선에 연결된 인접 정점이 '연결된 노드 집합'에 들어 있지 않으면, 해당 간선을 선택하고, 해당 간선 정보를 '최소 신장 트리(mst)'에 삽입

 

위에서 작성한 구현 알고리즘을 그대로 작성해 주시면 됩니다.

connected_node에 n2 노드를 추가하고 mst에 해당 간선 정보를 저장합니다. 

 

그리고 나서 새롭게 연결된 n2에 연결되어 있는 간선들의 정보를 candidate_edge_list에 저장해 주는 겁니다. 

 

해당 간선에 연결된 인접 정점의 간선들 중, '연결된 노드 집합(connected_nodes)' 에 없는 노드와 연결된 간선들만 간선 리스트(candidate_edge_list) 에 삽입 

 

앞서 adjacent_edges는 모든 쌍방의 간선 정보들을 저장하고 있다고 말씀드렸죠?

여기서 노드 n2에서 시작하는 모든 간선들을 조회했을 때 각 간선들의 목적지(이어주는 노드) 노드가 이미 connected_nodes에 저장이 되어있다면? 무시해버리면 됩니다! 어차피 안쓸거니까요 

저장이 안되어있는 것들만 후보 간선 리스트, candidate_edge_list에 넣어주시면 되는겁니다! 

 

그런식으로 모든 후보 간선들이 다 제거가 되고나면 (while candidate_edge_list) 이제 최소신장트리가 완성되었다는

뜻이며 이 최종결과값을 반환 해 주기만 하면 알고리즘은 끝이 납니다

 

결과값 :

prim ('A', myedges)

[(5, 'A', 'D'), (6, 'D', 'F'), (7, 'A', 'B'), (7, 'B', 'E'), (5, 'E', 'C'), (9, 'E', 'G')]

 

 

이 프림알고리즘의 시간 복잡도는 최악의 경우, while 구문에서 모든 간선에 대해 반복하고,

최소 힙 구조를 사용하므로 O(ElogE) 시간 복잡도를 가집니다. 

 

이를 개선하여 O(ElogV)의 시간 복잡도를 가진 프림알고리즘도 있긴 하지만 이는 기회가된다면 다음에 다루도록 하겠습니다. 이로서 프림알고리즘에 대한 포스트를 마치겠습니다.