prim.py

"""Prim's Algorithm.

Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm.

Details: https://en.wikipedia.org/wiki/Prim%27s_algorithm
"""

importheapqashq
importmath
fromcollections.abcimportIterator


classVertex:
"""Class Vertex."""

def__init__(self, id_):
"""
 Arguments:
 id - input an id to identify the vertex
 Attributes:
 neighbors - a list of the vertices it is linked to
 edges - a dict to store the edges's weight
 """
self.id=str(id_)
self.key=None
self.pi=None
self.neighbors= []
self.edges= {} # {vertex:distance}

def__lt__(self, other):
"""Comparison rule to < operator."""
returnself.key<other.key

def__repr__(self):
"""Return the vertex id."""
returnself.id

defadd_neighbor(self, vertex):
"""Add a pointer to a vertex at neighbor's list."""
self.neighbors.append(vertex)

defadd_edge(self, vertex, weight):
"""Destination vertex and weight."""
self.edges[vertex.id] =weight


defconnect(graph, a, b, edge):
# add the neighbors:
graph[a-1].add_neighbor(graph[b-1])
graph[b-1].add_neighbor(graph[a-1])
# add the edges:
graph[a-1].add_edge(graph[b-1], edge)
graph[b-1].add_edge(graph[a-1], edge)


defprim(graph: list, root: Vertex) ->list:
"""Prim's Algorithm.

 Runtime:
 O(mn) with `m` edges and `n` vertices

 Return:
 List with the edges of a Minimum Spanning Tree

 Usage:
 prim(graph, graph[0])
 """
a= []
foruingraph:
u.key=math.inf
u.pi=None
root.key=0
q=graph[:]
whileq:
u=min(q)
q.remove(u)
forvinu.neighbors:
if (vinq) and (u.edges[v.id] <v.key):
v.pi=u
v.key=u.edges[v.id]
foriinrange(1, len(graph)):
a.append((int(graph[i].id) +1, int(graph[i].pi.id) +1))
returna


defprim_heap(graph: list, root: Vertex) ->Iterator[tuple]:
"""Prim's Algorithm with min heap.

 Runtime:
 O((m + n)log n) with `m` edges and `n` vertices

 Yield:
 Edges of a Minimum Spanning Tree

 Usage:
 prim(graph, graph[0])
 """
foruingraph:
u.key=math.inf
u.pi=None
root.key=0

h=list(graph)
hq.heapify(h)

whileh:
u=hq.heappop(h)
forvinu.neighbors:
if (vinh) and (u.edges[v.id] <v.key):
v.pi=u
v.key=u.edges[v.id]
hq.heapify(h)

foriinrange(1, len(graph)):
yield (int(graph[i].id) +1, int(graph[i].pi.id) +1)


deftest_vector() ->None:
"""
 # Creates a list to store x vertices.
 >>> x = 5
 >>> G = [Vertex(n) for n in range(x)]

 >>> connect(G, 1, 2, 15)
 >>> connect(G, 1, 3, 12)
 >>> connect(G, 2, 4, 13)
 >>> connect(G, 2, 5, 5)
 >>> connect(G, 3, 2, 6)
 >>> connect(G, 3, 4, 6)
 >>> connect(G, 0, 0, 0) # Generate the minimum spanning tree:
 >>> G_heap = G[:]
 >>> MST = prim(G, G[0])
 >>> MST_heap = prim_heap(G, G[0])
 >>> for i in MST:
 ... print(i)
 (2, 3)
 (3, 1)
 (4, 3)
 (5, 2)
 >>> for i in MST_heap:
 ... print(i)
 (2, 3)
 (3, 1)
 (4, 3)
 (5, 2)
 """


if__name__=="__main__":
importdoctest

doctest.testmod()