Skip to content

Commit 724b7d2

Browse files
kylepwpoyea
authored andcommitted
Add Prim's algorithm with min heap (TheAlgorithms#1704)
1 parent 3042702 commit 724b7d2

File tree

1 file changed

+57
-8
lines changed

1 file changed

+57
-8
lines changed

graphs/prim.py

+57-8
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
1-
"""
2-
Prim's Algorithm.
1+
"""Prim's Algorithm.
2+
3+
Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm.
34
4-
Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm
5+
Details: https://en.wikipedia.org/wiki/Prim%27s_algorithm
56
"""
67

8+
import heapq as hq
79
import math
10+
from typing import Iterator
811

912

1013
class Vertex:
@@ -50,11 +53,17 @@ def connect(graph, a, b, edge):
5053
graph[b - 1].add_edge(graph[a - 1], edge)
5154

5255

53-
def prim(graph, root):
54-
"""
55-
Prim's Algorithm.
56-
Return a list with the edges of a Minimum Spanning Tree
57-
prim(graph, graph[0])
56+
def prim(graph: list, root: Vertex) -> list:
57+
"""Prim's Algorithm.
58+
59+
Runtime:
60+
O(mn) with `m` edges and `n` vertices
61+
62+
Return:
63+
List with the edges of a Minimum Spanning Tree
64+
65+
Usage:
66+
prim(graph, graph[0])
5867
"""
5968
a = []
6069
for u in graph:
@@ -74,6 +83,38 @@ def prim(graph, root):
7483
return a
7584

7685

86+
def prim_heap(graph: list, root: Vertex) -> Iterator[tuple]:
87+
"""Prim's Algorithm with min heap.
88+
89+
Runtime:
90+
O((m + n)log n) with `m` edges and `n` vertices
91+
92+
Yield:
93+
Edges of a Minimum Spanning Tree
94+
95+
Usage:
96+
prim(graph, graph[0])
97+
"""
98+
for u in graph:
99+
u.key = math.inf
100+
u.pi = None
101+
root.key = 0
102+
103+
h = [v for v in graph]
104+
hq.heapify(h)
105+
106+
while h:
107+
u = hq.heappop(h)
108+
for v in u.neighbors:
109+
if (v in h) and (u.edges[v.id] < v.key):
110+
v.pi = u
111+
v.key = u.edges[v.id]
112+
hq.heapify(h)
113+
114+
for i in range(1, len(graph)):
115+
yield (int(graph[i].id) + 1, int(graph[i].pi.id) + 1)
116+
117+
77118
def test_vector() -> None:
78119
"""
79120
# Creates a list to store x vertices.
@@ -87,13 +128,21 @@ def test_vector() -> None:
87128
>>> connect(G, 3, 2, 6)
88129
>>> connect(G, 3, 4, 6)
89130
>>> connect(G, 0, 0, 0) # Generate the minimum spanning tree:
131+
>>> G_heap = G[:]
90132
>>> MST = prim(G, G[0])
133+
>>> MST_heap = prim_heap(G, G[0])
91134
>>> for i in MST:
92135
... print(i)
93136
(2, 3)
94137
(3, 1)
95138
(4, 3)
96139
(5, 2)
140+
>>> for i in MST_heap:
141+
... print(i)
142+
(2, 3)
143+
(3, 1)
144+
(4, 3)
145+
(5, 2)
97146
"""
98147

99148

0 commit comments

Comments
 (0)