"""MinimumSpanningTree.py Kruskal's algorithm for minimum spanning trees. D. Eppstein, April 2006. """ import unittest from UnionFind import UnionFind from Graphs import isUndirected def MinimumSpanningTree(G): """ Return the minimum spanning tree of an undirected graph G. G should be represented in such a way that iter(G) lists its vertices, iter(G[u]) lists the neighbors of u, G[u][v] gives the length of edge u,v, and G[u][v] should always equal G[v][u]. The tree is returned as a list of edges. """ if not isUndirected(G): raise ValueError("MinimumSpanningTree: input is not undirected") for u in G: for v in G[u]: if G[u][v] != G[v][u]: raise ValueError("MinimumSpanningTree: asymmetric weights") # Kruskal's algorithm: sort edges by weight, and add them one at a time. # We use Kruskal's algorithm, first because it is very simple to # implement once UnionFind exists, and second, because the only slow # part (the sort) is sped up by being built in to Python. subtrees = UnionFind() tree = [] for W,u,v in sorted((G[u][v],u,v) for u in G for v in G[u]): if subtrees[u] != subtrees[v]: tree.append((u,v)) subtrees.union(u,v) return tree # If run standalone, perform unit tests class MSTTest(unittest.TestCase): def testMST(self): """Check that MinimumSpanningTree returns the correct answer.""" G = {0:{1:11,2:13,3:12},1:{0:11,3:14},2:{0:13,3:10},3:{0:12,1:14,2:10}} T = [(2,3),(0,1),(0,3)] for e,f in zip(MinimumSpanningTree(G),T): self.assertEqual(min(e),min(f)) self.assertEqual(max(e),max(f)) if __name__ == "__main__": unittest.main()