Create prim.py (TheAlgorithms#397)

Author: dunderbruno

Graphs/prim.py

@@ -0,0 +1,82 @@
+"""
+Prim's Algorithm.
+
+Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm
+
+Create a list to store x the vertices.
+G = [vertex(n) for n in range(x)]
+
+For each vertex in G, add the neighbors:
+G[x].addNeighbor(G[y])
+G[y].addNeighbor(G[x])
+
+For each vertex in G, add the edges:
+G[x].addEdge(G[y], w)
+G[y].addEdge(G[x], w)
+
+To solve run:
+MST = prim(G, G[0])
+"""
+
+importmath
+
+
+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."""
+return (self.key<other.key)
+
+def__repr__(self):
+"""Return the vertex id."""
+returnself.id
+
+defaddNeighbor(self, vertex):
+"""Add a pointer to a vertex at neighbor's list."""
+self.neighbors.append(vertex)
+
+defaddEdge(self, vertex, weight):
+"""Destination vertex and weight."""
+self.edges[vertex.id] =weight
+
+
+defprim(graph, root):
+"""
+ Prim's Algorithm.
+
+ Return a list with the edges of a Minimum Spanning Tree
+
+ 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([graph[i].id, graph[i].pi.id])
+return(A)