Krushkal's Algo for MST. Python Implementation

krshkals_algo.py

# Krushkal's Algorithm for Minimum Spanning Tree (MST)
# Python Implementation

from collections import defaultdict
from heapq import heappop, heappush


class Graph:
    def __init__(self):
        self.edge_heap = list()
        self.mst_edges = list()
        self.mst_weight = 0

        self.parent = defaultdict(lambda: None)
        self.rank = defaultdict(lambda: 1)

    def find(self, element):
        parent = self.parent[element]
        real_element = element

        while parent is not None:
            element = parent
            parent = self.parent[element]

        # Compress path
        if self.parent[real_element] is not None and self.parent[real_element] != element:
            self.parent[real_element] = element

        return element

    def union(self, x, y):
        parent_x = self.find(x)
        parent_y = self.find(y)

        if parent_x == parent_y:
            return

        rank_parent_x = self.rank[parent_x]
        rank_parent_y = self.rank[parent_y]

        if rank_parent_x > rank_parent_y:
            self.parent[parent_y] = parent_x
            self.rank[parent_x] += self.rank[parent_y]
        else:
            self.parent[parent_x] = parent_y
            self.rank[parent_y] += self.rank[parent_x]

    def get_mst_krushkals(self, edges):

        # Push all edges to the min Heap
        for u, v, w in edges:
            heappush(self.edge_heap, (int(w), (u, v)))

        while self.edge_heap:
            weight, uv = heappop(self.edge_heap)
            u, v = uv

            # Ensure that they are not forming a cycle. 
            # Use union find to achieve near constant time
            if self.find(u) == self.find(v):
                continue

            self.union(u, v)
            self.mst_weight += weight
            self.mst_edges.append((u, v, weight))

        return self.mst_weight


if __name__ == '__main__':
    edges = ["AB2", "AC3", "BC4", "DC5", "AD3", "DF7", "CF6", "CE1", "BE3", "FG9", "FE8"]

    graph = Graph()
    mst_weight = graph.get_mst_krushkals(edges=edges)
    print(f'MST Weight -> {mst_weight}')