[파이썬] 프림 알고리즘의 동작 원리와 예시

04 Sep 2023

python

프림 알고리즘은 그래프의 최소 신장 트리(Minimum Spanning Tree, MST)를 찾는 알고리즘 중 하나입니다. 이 알고리즘은 시작 정점에서부터 출발하여 신장 트리를 구성하는 과정에서 최소 비용의 간선을 선택하여 트리에 추가합니다.

프림 알고리즘의 동작 원리는 다음과 같습니다:

임의의 정점을 선택하여 트리에 포함시킵니다.
현재 트리에 속한 정점과 트리에 속하지 않은 정점 사이의 최소 비용의 간선을 선택합니다.
선택한 간선의 도착 정점을 새로운 정점으로 트리에 포함시킵니다.
모든 정점이 트리에 포함될 때까지 위 과정을 반복합니다.

이제 파이썬으로 프림 알고리즘을 구현한 코드를 살펴보겠습니다. 아래 예시는 주어진 그래프에서 최소 신장 트리를 구하는 프림 알고리즘의 예시입니다.

from heapq import heappop, heappush

def prim(graph):
    num_vertices = len(graph)
    
    # 임의의 정점으로 시작합니다 (여기서는 0번 정점을 선택합니다)
    start_vertex = 0
    
    # 트리의 크기를 저장할 리스트를 초기화합니다
    tree_size = [float('inf')] * num_vertices
    
    # 최소 신장 트리를 저장할 리스트를 초기화합니다
    mst = [None] * num_vertices
    
    # 힙 (min-heap)에 시작 정점을 추가합니다
    pq = [(0, start_vertex)]
    tree_size[start_vertex] = 0
    
    while pq:
        # 최소 비용의 간선을 선택합니다
        weight, current_vertex = heappop(pq)
        
        # 현재 간선의 도착 정점이 이미 트리에 속해있다면 스킵합니다
        if mst[current_vertex] is not None:
            continue
        
        # 도착 정점을 트리에 추가합니다
        mst[current_vertex] = (weight, current_vertex)
        
        # 현재 정점과 연결된 간선을 검사합니다
        for neighbor, edge_weight in graph[current_vertex]:
            # 트리에 속하지 않은 정점이라면 비용을 업데이트하고 힙에 추가합니다
            if mst[neighbor] is None and tree_size[neighbor] > edge_weight:
                tree_size[neighbor] = edge_weight
                heappush(pq, (edge_weight, neighbor))
    
    return mst

위 코드는 그래프를 인접 리스트로 표현하고, 힙을 사용하여 최소 비용의 간선을 선택하는 방식으로 프림 알고리즘을 구현한 예시입니다. 프림 알고리즘을 실행하면 시작 정점부터 출발하여 최소 신장 트리를 구성하게 됩니다.

프림 알고리즘은 정점의 개수가 V, 간선의 개수가 E일 때, O(E log V)의 시간 복잡도를 갖습니다. 따라서 아주 큰 그래프에도 효율적으로 적용할 수 있는 알고리즘입니다.