Added Kruskals MST

prakhar1989 · prakhar1989 · commit 59ab6e393907 · 2012-12-12T02:18:25.000+05:30
diff --git a/README.md b/README.md
@@ -2,7 +2,6 @@ Algos in Python
 ======
 
 Implementation for the algorithms and datastructures taught in [CS161](http://openclassroom.stanford.edu/MainFolder/CoursePage.php?course=IntroToAlgorithms), in Python
-Includes test code - doctests for few and full unittests for others
 
 Completed
 ---
diff --git a/graphs/graph_algorithms.py b/graphs/graph_algorithms.py
@@ -1,8 +1,7 @@
 from collections import deque
-from graph import graph
-from digraph import digraph
 from copy import deepcopy
 import heapq
+from union_find.unionfind import UnionFind
 
 def BFS(gr, s):
     """ Breadth first search 
@@ -219,3 +218,17 @@ def compute_key(gr, n, nodes_explored):
             w = gr.get_edge_weight((n, v))
             if w < min: min = w
     return min
+
+def kruskal_MST(gr):
+    """ computes minimum cost spanning tree in a undirected, 
+    connected graph using Kruskal's MST. Uses union-find data structure
+    for running times of O(mlogn) """
+    sorted_edges = sorted(gr.get_edge_weights())
+    uf = UnionFind()
+    min_cost = 0
+    for (w, (u, v)) in sorted_edges:
+        if (not uf.get_leader(u) and not uf.get_leader(v)) \
+                or (uf.get_leader(u) != uf.get_leader(v)):
+            uf.insert(u, v)
+            min_cost += w
+    return min_cost
diff --git a/tests/assign.py b/tests/assign.py
@@ -1,8 +1,7 @@
 import os, sys
 sys.path.append(os.path.join(os.getcwd(), os.path.pardir))
-
-from graphs import digraph
-from heaps import minheap
+from graphs.graph import graph
+from union_find.unionfind import UnionFind
 
 lines = [l for l in open("edges.txt").readlines()]
 lines = lines[1:]
@@ -16,10 +15,15 @@
         gr.add_node(v)
     gr.add_edge((u, v), int(w))
 
-
 def kruskal_MST(gr):
     sorted_edges = sorted(gr.get_edge_weights())
-    temp_gr = {}
+    uf = UnionFind()
     min_cost = 0
     for (w, (u, v)) in sorted_edges:
-        temp_gr[u] = v
+        if (not uf.get_leader(u) and not uf.get_leader(v)) \
+                or (uf.get_leader(u) != uf.get_leader(v)):
+            uf.insert(u, v)
+            min_cost += w
+    return min_cost
+
+print kruskal_MST(gr)
diff --git a/tests/graph_algorithms_test.py b/tests/graph_algorithms_test.py
@@ -115,5 +115,20 @@ def test_prims_minimum_spanning_tree(self):
         min_cost = minimum_spanning_tree(gr)
         self.assertEqual(min_cost, 39)
 
+    def test_kruskals_minimum_spanning_tree(self):
+        lines = [l for l in open("tests/edges.txt")]
+        lines = lines[1:]
+        edges = (l.split() for l in lines)
+        gr = graph()
+        for (u, v, w) in edges:
+            if u not in gr.nodes():
+                gr.add_node(u)
+            if v not in gr.nodes():
+                gr.add_node(v)
+            gr.add_edge( (u, v), int(w) )
+
+        min_cost = kruskal_MST(gr)
+        self.assertEqual(min_cost, 39)
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/union_find/unionfind.py b/union_find/unionfind.py
@@ -51,7 +51,7 @@ def make_union(self, leadera, leaderb):
         groupb = self.group[leaderb]
         if len(groupa) < len(groupb):
             # swap a and b if a is a smaller set
-            a, leadera, groupa, b, leaderb, groupb = b, leaderb, groupb, a, leadera, groupa
+            leadera, groupa, leaderb, groupb = leaderb, groupb, leadera, groupa
         groupa |= groupb # taking union of a with b
         del self.group[leaderb] # delete b
         for k in groupb:
diff --git a/union_find/unionfind_test.py b/union_find/unionfind_test.py
