Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 53 additions & 0 deletions dsu.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
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 = self.find(x)
root_y = 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

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 sort_edges(self):
self.graph.sort(key=lambda item: item[2])

def kruskal_mst(self):
dsu = DSU(self.V)
result = []
i, e = 0, 0

self.sort_edges()

while e < self.V - 1:
u, v, w = self.graph[i]
i += 1

if dsu.find(u) != dsu.find(v):
result.append([u, v, w])
dsu.union(u, v)
e += 1

return result, sum([weight for u, v, weight in result])
8 changes: 8 additions & 0 deletions graph.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
```python
import heapq

class Edge:
def __init__(self, u, v, w):
self.u = u
self.v = v
self.w = w