Implement Kruskal's Minimum Spanning Tree Algorithm

dsu.py

@@ -0,0 +1,37 @@
+class DSU:
+    def __init__(self, n):
+        self.parent = list(range(n))
+        self.rank = [0] * n
+
+    def find(self, x):
+        if self.parent[x] != x:
+            self.parent[x] = self.find(self.parent[x])
+        return self.parent[x]
+
+    def union(self, x, y):
+        root_x, root_y = self.find(x), self.find(y)
+        if root_x == root_y:
+            return
+        if self.rank[root_x] > self.rank[root_y]:
+            self.parent[root_y] = root_x
+        elif self.rank[root_x] < self.rank[root_y]:
+            self.parent[root_x] = root_y
+        else:
+            self.parent[root_y] = root_x
+            self.rank[root_x] += 1
+
+def kruskal_mst(graph):
+    v = len(graph.edges)
+    dsu = DSU(v)
+    edges = sorted(graph.edges, key=lambda edge: edge.weight)
+    mst_edges = []
+    total_weight = 0
+
+    for edge in edges:
+        x, y = dsu.find(edge.src), dsu.find(edge.dest)
+        if x != y:
+            mst_edges.append(edge)
+            total_weight += edge.weight
+            dsu.union(x, y)
+
+    return mst_edges, total_weight

graph.py

@@ -0,0 +1,64 @@
+import heapq
+
+class Graph:
+    def __init__(self, vertices):
+        self.V = vertices
+        self.graph = []
+
+    def add_edge(self, u, v, w):
+        self.graph.append([u, v, w])
+
+    def find(self, parent, i):
+        if parent[i] == i:
+            return i
+        return self.find(parent, parent[i])
+
+    def union(self, parent, rank, x, y):
+        xroot = self.find(parent, x)
+        yroot = self.find(parent, y)
+
+        if rank[xroot] < rank[yroot]:
+            parent[xroot] = yroot
+        elif rank[xroot] > rank[yroot]:
+            parent[yroot] = xroot
+        else:
+            parent[yroot] = xroot
+            rank[xroot] += 1
+
+    def kruskal_mst(self):
+        result = []
+        i, e = 0, 0
+        self.graph = sorted(self.graph, key=lambda item: item[2])
+        parent = []
+        rank = []
+
+        for node in range(self.V):
+            parent.append(node)
+            rank.append(0)
+
+        while e < self.V - 1:
+            u, v, w = self.graph[i]
+            i = i + 1
+            x = self.find(parent, u)
+            y = self.find(parent, v)
+
+            if x != y:
+                e = e + 1
+                result.append([u, v, w])
+                self.union(parent, rank, x, y)
+
+        return result, sum([edge[2] for edge in result])
+
+if __name__ == "__main__":
+    g = Graph(4)
+    g.add_edge(0, 1, 10)
+    g.add_edge(1, 2, 15)
+    g.add_edge(2, 0, 5)
+    g.add_edge(1, 3, 20)
+    g.add_edge(3, 2, 25)
+
+    mst, weight = g.kruskal_mst()
+    print("Edges in the MST are:")
+    for edge in mst:
+        print(f"{edge[0]}, {edge[1]}")
+    print(f"Minimum weight: {weight}")

test_kruskal.py

@@ -0,0 +1,53 @@
+import heapq
+
+def initialize_sets(n):
+    """Initialize n disjoint sets."""
+    return [i for i in range(n)]
+
+def find_set(rep, i):
+    """Find the set that i belongs to."""
+    if rep[i] == i:
+        return i
+    return find_set(rep, rep[i])
+
+def union_sets(rep, rank, x, y):
+    """Merge the sets containing x and y."""
+    xroot = find_set(rep, x)
+    yroot = find_set(rep, y)
+
+    if rank[xroot] < rank[yroot]:
+        rep[xroot] = yroot
+    elif rank[xroot] > rank[yroot]:
+        rep[yroot] = xroot
+    else:
+        rep[yroot] = xroot
+        rank[xroot] += 1
+
+def kruskal(graph):
+    """Find the minimum spanning tree using Kruskal's algorithm."""
+    v = graph.vertices()
+    e = graph.edges()
+
+    # Initialize the Disjoint Set Union (DSU) data structure
+    rep = initialize_sets(v)
+    rank = [0 for _ in range(v)]
+
+    # Sort the edges by weight
+    edges = sorted(e, key=lambda edge: edge.weight)
+
+    # Initialize the visited array and the MST
+    visited = [False for _ in range(e)]
+    mst_weight = 0
+
+    # Iterate through the sorted edges
+    for edge in edges:
+        x = find_set(rep, edge.start)
+        y = find_set(rep, edge.end)
+
+        # If the edge doesn't form a cycle, add it to the MST
+        if x != y:
+            union_sets(rep, rank, x, y)
+            visited[edge.index] = True
+            mst_weight += edge.weight
+
+    return mst_weight, visited