[Python] 최소신장트리(Minimum Spanning Trees, MST) 구현하기 by Prim

이번 글에서는 파이썬 언어를 사용하여 Prim방식 기반 최소신장트리를 구현한다.

1. 문제

가중치가 있는 무방향성 그래프가 주어졌을 때, 최소 신장 트리(MST, Minimum Spanning Tree)의 가중치의 합을 구하라

2. 파이썬 구현

기본적으로 주어지는 정보

import heapq

node_n, edge_n = 4,6
graph_info = [(0,1,6),(0,2,7),(2,3,2),(1,2,5),(1,3,3),(0,3,9)]
start_node = 0 # 탐색을 시작할 노드

주어진 그래프로부터 인접 딕셔너리 생성(인접 리스트로 해도 상관없음)

무방향 그래프이기때문에 node_a,node_b의 간선정보를 서로 넣어준다

# 인접 딕셔너리 생성
adj_dic=dict([(n,[])for n in range(node_n)])

# 그래프 정보를 읽어와서 인접 딕셔너리 작성
for node_a, node_b, edge_cost in graph_info:
    # 무방향성 그래프이기에 양쪽 다 넣어줌
    adj_dic[node_a].append((node_b,edge_cost))
    adj_dic[node_b].append((node_a,edge_cost))

방문기록을 저장 할 수 있는 리스트 생성(딕셔너리로 해도 가능)

# 방문기록용 리스트 생성
visit = [False for _ in range(node_n)]
visit[start_node] = True # 초기 노드 방문처리

우선순위 큐 생성 및 탐색 시작노드의 방문가능 간선정보로 초기화

# 우선순위 큐로 사용할 리스트 선언
que = []
# 탐색 시작노드에서 갈 수 있는 노드들 큐에 push
for node_b,cost in adj_dic[start_node]: #init_node가 node_a 역할
     heapq.heappush(que,(cost,node_b)) # cost가 우선순위 역할이기에 순서 변경

최소 간선길이로 모든 노드를 연결하는 문제이기에, 가중치(연결된 간선들의 길이)를 합할 수 있는 변수 선언

(연결된 간선 정보를 반환해야하면, 리스트로 만들어서 간선정보를 받아야함)

w_sum = 0 # 가중치의 합 // 추후 정답으로 반환

반복문을 통해서 우선순위 큐에서 가장 가중치가 높은 값을 pop함

pop된 노드가 방문된 기록이 없다면 방문하고 기록

방문한 노드에서 갈 수 있는 간선 정보들을 우선순위큐에 추가로 삽입

while que: # 큐가 비어있으면 반복 종료
    now_cost,now_nodeB = heapq.heappop(que) # 가장 우선순위가 높은것 pop
    if not visit[now_nodeB]: # pop된 노드가 방문된적 있는지 검사
        visit[now_nodeB] = True # 없다면 방문하고 방문표시
        w_sum += now_cost # 해당 노드로 가는 간선거리 추가
        for node_b,cost in adj_dic[now_nodeB]: # 해당 노드로부터 갈 수 있는 노드 큐에 삽입
            heapq.heappush(que,(cost,node_b))

3. 전체 코드

import heapq

node_n, edge_n = 4,6
graph_info = [(0,1,6),(0,2,7),(2,3,2),(1,2,5),(1,3,3),(0,3,9)]
start_node = 0 # 탐색을 시작할 노드

# 인접 딕셔너리 생성
adj_dic=dict([(n,[])for n in range(node_n)])

# 그래프 정보를 읽어와서 인접 딕셔너리 작성
for node_a, node_b, edge_cost in graph_info:
    # 무방향성 그래프이기에 양쪽 다 넣어줌
    adj_dic[node_a].append((node_b,edge_cost))
    adj_dic[node_b].append((node_a,edge_cost))

# 방문기록용 리스트 생성
visit = [False for _ in range(node_n)]
visit[start_node] = True # 초기 노드 방문처리
    
# 우선순위 큐로 사용할 리스트 선언
que = []
# 탐색 시작노드에서 갈 수 있는 노드들 큐에 push
for node_b,cost in adj_dic[start_node]: #init_node가 node_a 역할
     heapq.heappush(que,(cost,node_b)) # cost가 우선순위 역할이기에 순서 변경

w_sum = 0 # 가중치의 합 // 추후 정답으로 반환

while que: # 큐가 비어있으면 반복 종료
    now_cost,now_nodeB = heapq.heappop(que) # 가장 우선순위가 높은것 pop
    if not visit[now_nodeB]: # pop된 노드가 방문된적 있는지 검사
        visit[now_nodeB] = True # 없다면 방문하고 방문표시
        w_sum += now_cost # 해당 노드로 가는 간선거리 추가
        for node_b,cost in adj_dic[now_nodeB]: # 해당 노드로부터 갈 수 있는 노드 큐에 삽입
            heapq.heappush(que,(cost,node_b))
            
print(w_sum)

4. 코드 최적화

import heapq

node_n, edge_n = 4,6
graph_info = [(0,1,6),(0,2,7),(2,3,2),(1,2,5),(1,3,3),(0,3,9)]
start_node = 0 # 탐색을 시작할 노드

adj_dic=dict([(n,[])for n in range(node_n)])

for node_a, node_b, edge_cost in graph_info:
    adj_dic[node_a].append((edge_cost,node_b))
    adj_dic[node_b].append((edge_cost,node_a))


visit = [True]+[False for _ in range(node_n-1)]
que = adj_dic[start_node].copy()
w_sum = 0

while que: # 큐가 비어있으면 반복 종료
    now_cost,now_nodeB = heapq.heappop(que)
    if not visit[now_nodeB]:
        visit[now_nodeB] = True 
        w_sum += now_cost 
        for next_info in adj_dic[now_nodeB]:
            heapq.heappush(que,next_info)
            
print(w_sum)