Kruskals MST Algorithm This code computes the Minimum Spanning Tree of a given graph using Kruskals Algorithm.

It works successfully and I have provided test cases within the code.

I would like feedback on code efficiency (Choice of ds and functions/function size) and anything else to improve the code (Aside from pep 8 compliance, that is something I will be working to fix soon)

Thank you

class Graph:

  def __init__(self):

    self.graph = []
    self.edges = []
    self.parents = {}
    self.size = {}

  def addnode(self, key1, key2, edge=0):

    if key1 not in self.graph:
      self.graph.append(key1)
    if key2 not in self.graph:
      self.graph.append(key2)

    self.edges.append([key1, key2, edge])

  def getparent(self, vertex):

    while self.parents[vertex] != -1:
      vertex = self.parents[vertex]

    return vertex


  def union(self, vertex, neighbour):

    x_par = self.getparent(vertex)
    y_par = self.getparent(neighbour)
    if self.size[x_par] &lt;= self.size[y_par]:
      self.parents[x_par] = y_par
      self.size[y_par] += self.size[x_par]
    else:
      self.parents[y_par] = x_par
      self.size[x_par] += self.size[y_par]

  def kruskal(self):

    mst = []
    total = 0

    for key in self.graph:
      self.parents[key] = -1
      self.size[key] = 1

    self.edges.sort(key = lambda x:x[2])

    for edge in self.edges:
        vertex, neighbour, weight_edge = edge
        x_par = self.getparent(vertex)
        y_par = self.getparent(neighbour)
        if x_par != y_par:
          self.union(x_par, y_par)
          mst.append(edge)
          total += weight_edge

    print(mst)
    print('mst of weight:', total)

if __name__ == '__main__':
  g = Graph() 
  g.addnode(0, 1, 10) 
  g.addnode(0, 2, 6) 
  g.addnode(0, 3, 5) 
  g.addnode(1, 3, 15) 
  g.addnode(2, 3, 4) 
  g.kruskal()
Performance

The main cost of your implementation is in two areas, the sorting and the set memebership test.

When you do check for parents, it is generally advisable to update each of the parent of each of the ones you looked at, such that they would only be one step away from their parent the next time. Here is how it could look:

  def shortcut(self, vertex, parent):
    while self.parents[vertex] != -1:
      next = self.parents[vertex]
      self.parents[vertex] = parent
      vertex = next

  def getparent(self, vertex):
    start = vertex
    while self.parents[vertex] != -1:
      vertex = self.parents[vertex]

    self.shortcut(start, vertex)
    return vertex


It can be shown that this assymtotically reduces the cost of the set membership test in this kind of union based data-structure.

You can also cut these comparisons short when you have found the full minimal spanning tree, which would save you all the extra parent lookups for all the edges following in the sorted sequence. You can do this by checking whether you are done after adding an edge with len(mst) == len(self.graph) - 1, and then do a break or return the result, so the code (the loop in the kruskal function) would look like this:

    for edge in self.edges:
        vertex, neighbour, weight_edge = edge
        x_par = self.getparent(vertex)
        y_par = self.getparent(neighbour)
        if x_par != y_par:
          self.union(x_par, y_par)
          mst.append(edge)
          total += weight_edge
          if len(mst) == len(self.graph) - 1:
            break


For sorting you have two problems:


You specify a key function manually, which means the algorithm will have to call that key function.
You sort more of the problem than you really need to.


The first part comes from that sort may be implemented in a faster language (it looks to be so in my installation), which means that it generally runs faster if it never have to enter python code, which the manual key function breaks. If you however changes the order of attributes in how you store your edges, such that the weight is first, then you can use the default comparison, which will do lexiographic ordering, which in turn means you get it sorted primarily on the edge weights.

For the second part, the main problem is that you only need to search until you get the n-1 edges needed to finish the minimum spanning tree, so you only need this amount sorted. It would therefore be advisable to use a sorting scheme, such as heapsort, where you can stop half-way through and only need to effectively pay in performance for the amount you needed. The downside is that the heapq python module looks to be mainly implemented in python, which means its actual performance might be slower unless you really only need a very small subset (this would typically depend on the graph).

Convenience

There are a few areas in the code, which with some improvement would be more convinient/better future-proofed.

The main one here is the need for a simpler way to add several edges to a graph, or even construct a graph with edges already in it. Here is how I would do it:

  def __init__(self, edges=None):

    self.graph = []
    self.edges = []
    self.parents = {}
    self.size = {}
    self.add_edges(edges)

  def add_edges(self, edges=None):
    if edges is not None:
      for edge in edges:
        self.add_node(*edge)


The construction of your graph would then be this simple:

  g = Graph([
    (0, 1, 10) 
    (0, 2, 6) 
    (0, 3, 5) 
    (1, 3, 15) 
    (2, 3, 4) ])


In the union fuction, we could parameterize the if statement as follows:

    child, parent = (x_par, y_par) if self.size[x_par] &lt;= self.size[y_par] else (y_par, x_par)
    self.parents[child] = parent 
    self.size[parent] += self.size[child]


This would also seperate the orthogonal parts of the logic, which would make it easier to maintain.

Lastly it would be a good idea to have the kruskal method just return mts, total and then print outside, as this would allow actual use of the class in code outside of problems where this is the final solution.