diff --git a/dsu.py b/dsu.py new file mode 100644 index 0000000000..a21add72d8 --- /dev/null +++ b/dsu.py @@ -0,0 +1,33 @@ +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): + edge_list = sorted((weight, u, v) for u, v, weight in graph.edges(data='weight')) + dsu = DSU(graph.number_of_nodes()) + mst_weight = 0 + + for weight, u, v in edge_list: + if dsu.find(u) != dsu.find(v): + dsu.union(u, v) + mst_weight += weight + + return mst_weight \ No newline at end of file diff --git a/graph.py b/graph.py new file mode 100644 index 0000000000..af667887d0 --- /dev/null +++ b/graph.py @@ -0,0 +1,60 @@ +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]) + + # find set of an element i + def find(self, parent, i): + if parent[i] == i: + return i + return self.find(parent, parent[i]) + + # union of two sets of x and y + def union(self, parent, rank, x, y): + xroot = self.find(parent, x) + yroot = self.find(parent, y) + + # attach smaller rank tree under root of high rank tree + if rank[xroot] < rank[yroot]: + parent[xroot] = yroot + elif rank[xroot] > rank[yroot]: + parent[yroot] = xroot + else: + parent[yroot] = xroot + rank[xroot] += 1 + + # main function to construct MST using Kruskal's algorithm + def kruskal_mst(self): + result = [] + i, e = 0, 0 + + # sort all edges in non-decreasing order of their weight + self.graph = sorted(self.graph, key=lambda item: item[2]) + + parent = [] + rank = [] + + # create V subsets with single elements + for node in range(self.V): + parent.append(node) + rank.append(0) + + # number of edges to be taken is equal to V-1 + 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 including this edge does not cause cycle, include it in result + if x != y: + e = e + 1 + result.append([u, v, w]) + self.union(parent, rank, x, y) + + return result, sum([e[2] for e in result]) \ No newline at end of file diff --git a/kruskal.py b/kruskal.py new file mode 100644 index 0000000000..10f0009ad3 --- /dev/null +++ b/kruskal.py @@ -0,0 +1,8 @@ + ```python +import heapq + +class Edge: + def __init__(self, u, v, w): + self.u = u + self.v = v + self.w = w