Prim's Algorithm
This is an implementation of Prim's algorithm in Python. From Wikipedia:
- Initialize a tree with a single vertex, chosen arbitrarily from the graph.
- Grow the tree by one edge: of the edges that connect the tree to vertices not yet in the tree, find the minimum-weight edge, and transfer it to the tree.
- Repeat step 2 (until all vertices are in the tree).
My code
I have included the all the relevant sections for completeness but I want advice on the Prim's function inside the Graph class but feel free to comment on any part of the code!
I feel like the Prim's function can be improved as I have lots of conditionals that are similar but I don't know how to make it more Pythonic.
Any critique is welcome.
class Vertex:
def __init__(self, name):
self.name = name
def __str__(self):
return f"Vertex {self.name}"
class Edge:
def __init__(self, start, end, weight,directed):
self.start = start
self.end = end
self.weight = weight
self.directed = directed
def __str__(self):
return f"{self.start.name}{self.end.name}"
class Graph:
def __init__(self, v, e):
self.vertices = v
self.edges = e
def add_vertex(self, v):
""" Add vertex of type Vertex. """
self.vertices.append(v)
def total_weight(self):
""" Return total weight of all edges in graph."""
return sum(e.weight for e in self.edges)
def vertex_from_name(self, name):
""" Return vertex object given vertex name. """
return next((v for v in self.vertices if v.name == name), None)
def add_edge(self, start, end, weight,directed=False):
""" Add an edge connecting two vertices. Arguments can either be vertex name or vertex object. """
if isinstance(start, str):
start = self.vertex_from_name(start)
if isinstance(end, str):
end = self.vertex_from_name(end)
self.edges.append(Edge(start, end, weight,directed))
def add_edges(self,edges):
for edge in edges:
self.add_edge(edge[0],edge[1],edge[2])
def edge_on_vertex(self, v):
""" Return edges connected to given vertex v."""
return (e for e in self.edges if v in {e.start, e.end})
def connected_vertices(self, v):
""" Return the vertices connected to argument v."""
if isinstance(v, str):
v = self.vertex_from_name(v)
yield from (e.start for e in self.edges if e.end == v)
yield from (e.end for e in self.edges if e.start == v)
#Code to be reviewed
def Prims(self, **kwargs):
""" Return MST using Prim's algorithm. Optional argument is start vertex, defaults to first vertex. """
self.start = kwargs.get('start', self.vertices[0])
self.tree = Graph([], [])
self.tree.vertices.append(self.start)
while len(self.tree.vertices) != len(self.vertices):
self.connected = set([e for vert in self.tree.vertices for e in self.edge_on_vertex(vert)])
self.connected = sorted(list(self.connected), key=lambda x: x.weight)
for edge in self.connected:
if (edge.start not in self.tree.vertices) or (edge.end not in self.tree.vertices):
if edge.start in self.tree.vertices:
self.tree.add_vertex(edge.end)
else:
self.tree.add_vertex(edge.start)
self.tree.edges.append(edge)
break
return self.tree
if __name__ == "__main__":
v = [Vertex(x) for x in 'ABCDEF']
g = Graph(v, [])
g.add_edges((
("A", "B", 9),
("A", "C", 12),
("A", "D", 9),
("A", "E", 11),
("A", "F", 8),
("B", "C", 10),
("B", "F", 15),
("C", "D", 8),
("D", "E", 14),
("E", "F", 12),
))
print([str(e) for e in g.Prims().edges])
Code seem familiar? This is a followup to my earlier question about Kruskal's algorithm using the same module.
