빽 투더 기본기 [알고&자구 7편]. 프림 알고리즘

그래프에서 최소 신장 트리를 만드는 알고리즘 중, 크루스칼과 더불어 유명한 프림(Prim) 알고리즘을 잠깐 정리해볼까 한다.

글 대부분의 내용은 프로그래밍 자료 저장소님의 블로그와 Michael Sambol의 유튜브를 참고하였다.

 

그래프에서 최소 신장 트리 만들기

프림 알고리즘은, 크루스칼 알고리즘 같이 어떤 그래프가 주어져있을 때, 이 그래프에서 최소 신장 트리 (Minimum Spanning Tree) 를 만드는 알고리즘이다.

최소 신장 트리에 대한 이야기는 이전 글인 크루스칼 알고리즘에 보면 나와있다.

프림 알고리즘을 위해서는, 역시 먼저 그래프가 주어져야 한다. 그리고 이 그래프에서, 우리는 정점과 인접 행렬에 대한 정보가 주어져야 한다.

 

프림(Prim) 알고리즘

알고리즘

알고리즘을 간단하게 정리하면 다음과 같다.

1. 그래프에서 임의의 한 정점을 잡아서 새로운 그래프에 넣는다.
2. 이 점에 연결된 (간선, 정점) 을 인접 행렬을 통해 찾아서, 최소 힙에 넣는다.
3. 최소 힙에서 하나를 pop 하여 새로운 그래프에 넣는다.
4. 모든 정점을 방문할 때까지 2-3을 반복한다.

이해가 안되면… 위에 링크한 Michael Sambol의 유튜브를 보면 2분만에 이해할 수 있다.

아무튼 이 알고리즘을 한마디로 표현하면,

프림 알고리즘은 가장 가까운 정점을 하나씩 방문하며 최소 신장 트리를 만들어가는 일종의 그리디 알고리즘이다.

크루스칼은 가장 작은 간선 위주로 탐색해나가며 그래프(트리) 를 만드는 방법이라면,
프림은 임의의 점을 시작으로 연결된 간선을 탐색하나가며 그래프(트리)를 만드는 방법이라 할 수 있겠다.

 

코드

먼저 그래프가 주어져야 한다고 했다.
아래와 같은 그래프 상황이고, 정점과 인접행렬이 주어졌다고 하자.

출처 : 허민석님 유튜브

v = list("ABCDEF")
adjacent = {
    "A": {
        "B": 3,
        "C": 2,
        "D": 4,
    },
    "B": {
        "A": 3,
        "D": 2,
        "F": 5
    },
    "C": {
        "A": 2,
        "E": 1
    },
    "D": {
        "A": 4,
        "B": 2,
        "E": 1,
        "F": 3
    },
    "E": {
        "C": 1,
        "D": 1,
        "F": 2,
    },
    "F": {
        "B": 5,
        "D": 3,
        "E": 2
    }
}

이제 위 알고리즘 대로 코딩을 하면,

import heapq

min_cost = 0
min_edges = []
min_heap = []
not_visited = set(adjacent.keys())

# 시작에 연결된 간선들을 최소 힙에 넣어준다.
start = 'A'
for v, weight in adjacent[start].items():
    heapq.heappush(min_heap, (weight, (start, v)))
not_visited.remove(start)

# 모든 정점을 방문할 때 까지 아래를 반복한다.
while not_visited:

    # 최소 힙에서 간선을 빼온다.
    weight, (v1, v2) = heapq.heappop(min_heap)

    # 이어지는 간선에 다음 정점이 방문되지 않은 정점이면,
    if v2 in not_visited:

        not_visited.remove(v2)
        min_cost += weight
        min_edges.append((weight, (v1, v2)))

        # 그 정점에 연결된 간선들을 최소힙에 넣는다.
        for next_v, weight in adjacent[v2].items():
            if next_v in not_visited:
                heapq.heappush(min_heap, (weight, (v2, next_v)))

이를 출력하면 다음과 같이 최소 비용과 최소 신장 트리에 대한 정보를 알 수 있다.

print(min_cost)
print(min_edges)

# output
8
[(2, ('A', 'C')),
 (1, ('C', 'E')),
 (1, ('E', 'D')),
 (2, ('D', 'B')),
 (2, ('E', 'F'))]

프림 알고리즘의 시간복잡도를 생각하면,

1. 그래프에서 모든 간선을 결국 최소힙에 넣으므로, 최악의 경우 O(|e|) 번 반복문을 돌고,
2. 각 반복문 안에서, 최소힙과 관련된 연산은 |log V| 다.

따라서, 결과적으로 시간복잡도는 O(e log V) 라고 생각할 수 있다.