[Coding Test][Python][Shortest Path] Prim 알고리즘

 

 

Prim 알고리즘에 앞서 Kruskal 알고리즘을 보고 오시면 좋습니다.

[Coding Test][Python][Shortest Path] Kruskal 알고리즘 개념 및 예제, Union - Find 자료구조

 

1. Prim 알고리즘이란?

Prim 알고리즘은 Kruskal 알고리즘과 같이 Graph에서 Minimum Spanning Tree(MST)를 찾는 알고리즘입니다.

특징

• Kruskal 알고리즘은 Edge를 중심으로 Prim 알고리즘은 Vertex를 중심으로 MST 확장
• Greedy 알고리즘
• 보통 Priority Queue를 사용
• 시간 복잡도 : $O(E log V)$

2. Prim 알고리즘 동작 과정

1) 임의의 Vertex를 시작점으로 선택하여 MST에 추가

Prim 알고리즘은 아무 Vertex에서 출발하면 됩니다.

 

2) MST에 속한 Vertex들과 인접한 Vertex 중 가장 낮은 Weight를 가진 Edge와 연결된 Ver에 대해 Edge과 Vertex를 MST에 넣는다.

1번 Vertex에 연결된 Edge들 중 Weight가 가장 작은 Edge를 선택하고 MST에 넣습니다.

현재 MST를 구성하고 있는 Edge는 (1, 3) 만 존재합니다.

이 MST에 새로운 Vertex를 더하기 위한 방법은 Vertex 1번과 연결된 Edge 중 Weight가 가장 작은 Edge를 연결하는 것입니다.

Vertex 1에 대해서 모든 Edge를 확인하였으므로, Weight가 11로 더 작았던 Vertex 3에 대해서도 확인합니다.

이때 2- 3 Edge를 포함하지 않는 것은 Vertex 2가 이미 MST에 존재하기 때문입니다.

3) 2번 과정을 MST가 모든 Vertex를 포함할 때까지 반복

Prim알고리즘에 대한 구현은 Dijkstra 알고리즘과 동일합니다.
Dijkstra 알고리즘에서 현재 Vertex u에서 v로 간다고 할 때, $dist[v]$와 $dist[u] + edge[u][v]$를 비교하여 갱신해주면 됩니다.
Prim 알고리즘에서는 단순하게 $dist[v]$와 $edges[u][v]$를 비교하여 갱신해주면 됩ㄴ디ㅏ.

3. Prim 알고리즘 구현

1) Edge 및 Graph 초기화

Dijkstra 알고리즘에서는 Graph를 정의할때 (v, weight)로 Graph를 구성하였는데, Prim 알고리즘에서는 weight를 통해 Priority Queue를 사용하므로 (weight, v) 형태로 저장합니다.

from collections import defaultdict
import heapq

V, E = 6, 7
# Input을 주어진다고 가정
edges = [
    [1, 2, 12],
    [1, 3, 11],
    [2, 3, 13],
    [2, 5, 17],
    [3, 4, 15],
    [3, 6, 14],
    [4, 5, 19]
]

graph = defaultdict(list)
for edge in edges:
    u, v, weight = edge
    # Priority Queue를 위해 Weight를 [0]에 저장장
    graph[u].append((weight, v))
    graph[v].append((weight, v))

2) Prim 알고리즘 구현

Dijkstra 알고리즘과 매우 유사한 것을 확인할 수 있습니다.

def prim(graph, start):
    MST = []
    total_weight = 0
    pq = []
    visited = []

    visited.append(start)

    for weight, v in graph[start]:
        heapq.heappush(pq, (weight, start, v))

    while pq:
        weight, u, v = heapq.heappop(pq)

        if v not in visited:
            # 선택된 Vertex v를 MST에 추가가
            visited.append(v)
            MST.append((u, v, weight))
            total_weight += weight

            # Vertex v와 연결된 Edge 추가
            for next_weight, next_node in graph[v]:
                if next_node not in visited:
                    heapq.heappush(pq, (next_weight, v, next_node))

    return MST, total_weight

3) 결과확인

MST, total_weight = prim(graph, 1)
print(MST)
print(total_weight)

[(1, 3, 11), (1, 2, 12), (3, 6, 14), (3, 4, 15), (2, 5, 17)]
69