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
79import math
10+ from typing import Iterator
811
912
1013class 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+
77118def 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