Prim's Algorithm

Question

Prim's Algorithm

This is an implementation of Prim's algorithm in Python. From Wikipedia:

1. Initialize a tree with a single vertex, chosen arbitrarily from the graph.
2. 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.
3. 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.

Answer 1

Be consistent

• You named all your method functions with lowercase letters and separated the words with underscores... So maybe Prims should be named prims or prim_algorithm;

• In lists (argument lists, function calls, lists, etc) you always include a space after a comma (as you should), so do it always (I'm talking about Edge.__init__, Graph.add_egde and Graph.add_edges, for e.g.);

• For the simpler Graph methods, the one-liners, you have three lines: function definition, docstring, return statement; this is perfectly fine for simple functions. So you should also do it for Graph.vertex_from_name, instead of having a blank line between the docstring and the return.

Docstring conventions

For multiline docstrings, consider having the final """ in a separate line, all by itself, and then include a blank line between the multiline docstring and the body of the function.

Avoid long lines

Python's style guide suggests your lines don't exceed 81 characters (generally) and the black code formatter goes with 89, because "it's like highway speed limits, we won't bother you if you overdo it by a few km/h". But don't exceed that, no one wants to have to scroll right when reading someone else's code; for example, in your question, I have to scroll right to read some docstrings and some code in your functions.

Prims

As for the Prims function itself, I would reformat the body slightly:

def Prims(self, start=None):
    """ Return MST using Prim's algorithm.

    Optional argument `start` gives the start vertex (defaults to first vertex).
    """

    if start is None:
        start = self.vertices[0]
    self.tree = Graph([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

• I reformatted the docstring to keep the first line short and to the point. A short first docstring line is very helpful because many IDEs can show it if you hover the function name when using it elsewhere. If the docstring is long and/or contains irrelevant information, you won't be able to read what you needed to recall what your function does.

• I added an explicit keyword argument, instead of having you guess what you decided the starting vertex argument was named; I also gave it a default value of None, against which I then compare to see if I need to use the default value. If you want that argument to always be called as a keyword argument, you can use this syntax: def prims_algorithm(self, *, start=None):.

• If Graph takes a list of vertices when creating a Graph instance, why don't you initialize your tree instance already with the starting vertex, instead of appending it right after instantiating the Graph?

• I added two blank lines by the end of the function to make it easier to spot the return and to make it easier to spot the break; in particular the break statement was fairly hard to find and a quick read didn't reveal it. I like having these important keywords visible!