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.
3
4
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
5
6
"""
6
7
8
+ import heapq as hq
7
9
import math
10
+ from typing import Iterator
8
11
9
12
10
13
class Vertex :
@@ -50,11 +53,17 @@ def connect(graph, a, b, edge):
50
53
graph [b - 1 ].add_edge (graph [a - 1 ], edge )
51
54
52
55
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])
58
67
"""
59
68
a = []
60
69
for u in graph :
@@ -74,6 +83,38 @@ def prim(graph, root):
74
83
return a
75
84
76
85
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
+
77
118
def test_vector () -> None :
78
119
"""
79
120
# Creates a list to store x vertices.
@@ -87,13 +128,21 @@ def test_vector() -> None:
87
128
>>> connect(G, 3, 2, 6)
88
129
>>> connect(G, 3, 4, 6)
89
130
>>> connect(G, 0, 0, 0) # Generate the minimum spanning tree:
131
+ >>> G_heap = G[:]
90
132
>>> MST = prim(G, G[0])
133
+ >>> MST_heap = prim_heap(G, G[0])
91
134
>>> for i in MST:
92
135
... print(i)
93
136
(2, 3)
94
137
(3, 1)
95
138
(4, 3)
96
139
(5, 2)
140
+ >>> for i in MST_heap:
141
+ ... print(i)
142
+ (2, 3)
143
+ (3, 1)
144
+ (4, 3)
145
+ (5, 2)
97
146
"""
98
147
99
148
0 commit comments