Ass 2 Q1 union-by-height -> rank

1 files changed, 7 insertions, 4 deletions

diff --git a/ass2/q1/kruskals.py b/ass2/q1/kruskals.py
index f3c0703..6b73019 100644
--- a/ass2/q1/kruskals.py
+++ b/ass2/q1/kruskals.py
@@ -53,15 +53,18 @@ def sort_edges(edges):
 def find(groups, vertex):
     """
     Finds the root of the vertex's group according to a union-find data structure
+    Implements the find operation of union-by-rank
     """
-    while groups[vertex] >= 0:
-        vertex = groups[vertex]
-    return vertex
+    if groups[vertex] < 0:
+        return vertex
+    else:
+        groups[vertex] = find(groups, groups[vertex])
+        return find(groups, groups[vertex])
 
 
 def union(groups, u, v):
     """
-    Unions two vertices together according to a union-find data structure
+    Unions two vertices together (by height/rank) according to a union-find data structure
     """
     root_u, root_v = find(groups, u), find(groups, v)
     height_u, height_v = -groups[root_u], -groups[root_v]