Graphs and Kruskal's Algorithm - Python OOP This is my implementation of a graph and Kruskal's algorithm in Python. I wanted to design the graph module myself and I would like some feedback on the design. I tried to follow SOLID throughout. I am not sure if the separate Vertex object is wise, but I feel it could be useful as I expand this module. 

I had a copy of a flowchart of Kruskal's algorithm from a textbook (not my current course) and I decided to implement it, I am wondering how Pythonic my code is.

I have also programmed Prim's algorithm in the same file but I will split it over two questions.

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):
        self.start = start
        self.end = end
        self.weight = weight

    def __str__(self):
        return f"{self.start}{self.end}"

class Graph:
    def __init__(self, v, e):
        self.vertices = v
        self.edges = e

    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):
        """ 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))

    def edge_on_vertex(self, v):
        """ Return edges connected to given vertex v."""
        return [e for e in self.edges if (e.start == v) or (e.end == v)]

    def connected_vertices(self, v):
        """ Return the vertices connected to argument v."""
        if isinstance(v, str):
            v = self.vertex_from_name(v)

        return [e.start for e in self.edges if e.end == v] + [e.end for e in self.edges if e.start == v]

    def union(self, lst, e1, e2):
        """ Given a list of lists, merges e1 root list with e2 root list and returns merged list."""
        xroot, yroot = [], []
        # Find roots of both elements
        for i in lst:
            if e1 in i:
                xroot = i
            if e2 in i:
                yroot = i
        # Same root, cannot merge
        if xroot == yroot:
            return False
        xroot += yroot
        lst.remove(yroot)
        return lst

    def is_cycle(self):
        """ Return if the graph contains a cycle. """
        self.sets = [[v] for v in self.vertices]
        self._edges = sorted(self.edges, key=lambda x: x.weight)
        for e in self._edges:
            _temp = self.union(self.sets, e.start, e.end)
            if _temp == False:
                return True
            else:
                self.sets = _temp
        return False

    def Kruskals(self):
        """ Return MST using Kruskal's algorithm. """
        self.tree = Graph([], [])
        self.tree.vertices = self.vertices
        self.sorted_edges = sorted(self.edges, key=lambda x: x.weight)
        self.tree.edges.append(self.sorted_edges.pop(0))

        for edge in self.sorted_edges:
            self.tree.edges.append(edge)
            if self.tree.is_cycle():
                self.tree.edges.remove(edge)

        return self.tree

if __name__ == "__main__":
    v = [Vertex(x) for x in ["A", "B", "C", "D", "E", "F"]]

    g = Graph(v, [])

    g.add_edge("A", "B", 9)
    g.add_edge("A", "C", 12)
    g.add_edge("A", "D", 9)
    g.add_edge("A", "E", 11)
    g.add_edge("A", "F", 8)
    g.add_edge("B", "C", 10)
    g.add_edge("B", "F", 15)
    g.add_edge("C", "D", 8)
    g.add_edge("D", "E", 14)
    g.add_edge("E", "F", 12)

    print(g.Kruskals().edges)
Type hints

def __init__(self, start, end, weight):


can be

def __init__(self, start: Vertex, end: Vertex, weight: float):


depending on a few things, including the order of declaration of your classes, Vertex might need to be 'Vertex' here.

For another example, this

def vertex_from_name(self, name):


can turn into

def vertex_from_name(self, name: str) -&gt; Vertex:


Efficient lookups

To make this more efficient:

    return next((v for v in self.vertices if v.name == name), None)


Consider maintaining a string-to-Vertex dictionary to reduce this lookup from O(n) to O(1) in time.

Premature materialization

These:

    return [e for e in self.edges if (e.start == v) or (e.end == v)]


    return [e.start for e in self.edges if e.end == v] + [e.end for e in self.edges if e.start == v]


require that the entire results be stored to an in-memory list. To return the generator directly and reduce this memory requirement, the first one can be

    return (e for e in self.edges if v in {e.start, e.end})


and the second one can be

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)


Set-membership tests

This:

""" Given a list of lists, merges e1 root list with e2 root list and returns merged list."""


is probably better-expressed as accepting a list of set, not a list of lists. That will make these two tests:

        if e1 in i:
            xroot = i
        if e2 in i:
            yroot = i


faster. This:

    self.sets = [[v] for v in self.vertices]


would then become

    self.sets = [{v} for v in self.vertices]


Strings as iterables

This

v = [Vertex(x) for x in ["A", "B", "C", "D", "E", "F"]]


can be

v = [Vertex(x) for x in 'ABCDEF']


Convenience functions

Consider making a convenience function to turn this

g.add_edge("A", "B", 9)
g.add_edge("A", "C", 12)
g.add_edge("A", "D", 9)
g.add_edge("A", "E", 11)
g.add_edge("A", "F", 8)
g.add_edge("B", "C", 10)
g.add_edge("B", "F", 15)
g.add_edge("C", "D", 8)
g.add_edge("D", "E", 14)
g.add_edge("E", "F", 12)


into

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),
)