【python】并查集

转自：http://blog.csdn.net/rav009/article/details/12781899

 1 # -*- coding: UTF-8 -*-
 2 
 3 class unionfind:
 4     def __init__(self, groups):
 5         self.groups = groups
 6         self.items = []
 7         for g in groups:
 8             self.items += list(g)
 9         self.items = set(self.items)
10         self.parent = {}
11         self.rootdict = {}  # 记住每个root下节点的数量
12         for item in self.items:
13             self.rootdict[item] = 1
14             self.parent[item] = item
15 
16     def union(self, r1, r2):
17         rr1 = self.findroot(r1)
18         rr2 = self.findroot(r2)
19         cr1 = self.rootdict[rr1]
20         cr2 = self.rootdict[rr2]
21         if cr1 >= cr2:  # 将节点数量较小的树归并给节点数更大的树
22             self.parent[rr2] = rr1
23             self.rootdict.pop(rr2)
24             self.rootdict[rr1] = cr1 + cr2
25         else:
26             self.parent[rr1] = rr2
27             self.rootdict.pop(rr1)
28             self.rootdict[rr2] = cr1 + cr2
29 
30     def findroot(self, r):
31         """
32         可以通过压缩路径来优化算法,即遍历路径上的每个节点直接指向根节点
33         """
34         if r in self.rootdict.keys():
35             return r
36         else:
37             return self.findroot(self.parent[r])
38 
39     def createtree(self):
40         for g in self.groups:
41             if len(g) < 2:
42                 continue
43             else:
44                 for i in range(0, len(g) - 1):
45                     if self.findroot(g[i]) != self.findroot(g[i + 1]):  # 如果处于同一个集合的节点有不同的根节点，归并之
46                         self.union(g[i], g[i + 1])
47 
48     def printree(self):
49         rs = {}
50         for item in self.items:
51             root = self.findroot(item)
52             rs.setdefault(root, [])
53             rs[root] += [item]
54         for key in rs.keys():
55             print rs[key],
56 
57     def gettree(self):
58         rs = {}
59         group = []
60         for item in self.items:
61             root = self.findroot(item)
62             rs.setdefault(root, [])
63             rs[root] += [item]
64         for key in rs.keys():
65             group.append(rs[key])
66         return group
67 
68 
69 if __name__ == '__main__':
70     u = unionfind([['a', 'b', 'c'], ['b', 'd'], ['e', 'f'], ['g'], ['d', 'h', 'i']])
71     u.createtree()
72     u.printree()