Kruskals.py

class Graph(object):

    # g = [[vertex_1, vertex_2, weight]]
    # vertices = # of vertices in graph
    def __init__(self, vertices, g):
        self.g = g
        self.vertices = vertices

    def kruskal_mst(self):
        result = []

        self.g = sorted(self.g, key=lambda item: item[2])

        parent = []
        rank = []

        for node in range(self.vertices):
            parent.append(node)
            rank.append(0)

        sorted_edges = 0
        edges_taken = 0
        while edges_taken < self.vertices - 1:
            u, v, w = self.g[sorted_edges]
            sorted_edges += 1 

            v1 = self.find(parent, u)
            v2 = self.find(parent, v)

            # does not form a cycle
            if v1 != v2:
                edges_taken += 1
                result.append([u,v,w])
                self.union(parent, rank, v1, v2)

        return result

    def find(self, parent, vertex):
        if parent[vertex] == vertex:
            return vertex
        return self.find(parent, parent[vertex])
        
    def union(self, parent, rank, v1, v2):
        v1_root = self.find(parent, v1)
        v2_root = self.find(parent, v2)

        if rank[v1_root] < rank[v2_root]:
            parent[v1_root] = v2_root
        elif rank[v1_root] > rank[v2_root]:
            parent[v2_root] = v1_root
        else:
            parent[v2_root] = v1_root
            rank[v1_root] += 1

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