【python】并查集

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

 # -*- coding: UTF-8 -*-

 class unionfind:
     def __init__(self, groups):
         self.groups = groups
         self.items = []
         for g in groups:
             self.items += list(g)
         self.items = set(self.items)
         self.parent = {}
         self.rootdict = {}  # 记住每个root下节点的数量
         for item in self.items:
             self.rootdict[item] = 1
             self.parent[item] = item

     def union(self, r1, r2):
         rr1 = self.findroot(r1)
         rr2 = self.findroot(r2)
         cr1 = self.rootdict[rr1]
         cr2 = self.rootdict[rr2]
         if cr1 >= cr2:  # 将节点数量较小的树归并给节点数更大的树
             self.parent[rr2] = rr1
             self.rootdict.pop(rr2)
             self.rootdict[rr1] = cr1 + cr2
         else:
             self.parent[rr1] = rr2
             self.rootdict.pop(rr1)
             self.rootdict[rr2] = cr1 + cr2

     def findroot(self, r):
         """
         可以通过压缩路径来优化算法,即遍历路径上的每个节点直接指向根节点
         """
         if r in self.rootdict.keys():
             return r
         else:
             return self.findroot(self.parent[r])

     def createtree(self):
         for g in self.groups:
             if len(g) < 2:
                 continue
             else:
                 for i in range(0, len(g) - 1):
                     if self.findroot(g[i]) != self.findroot(g[i + 1]):  # 如果处于同一个集合的节点有不同的根节点，归并之
                         self.union(g[i], g[i + 1])

     def printree(self):
         rs = {}
         for item in self.items:
             root = self.findroot(item)
             rs.setdefault(root, [])
             rs[root] += [item]
         for key in rs.keys():
             print rs[key],

     def gettree(self):
         rs = {}
         group = []
         for item in self.items:
             root = self.findroot(item)
             rs.setdefault(root, [])
             rs[root] += [item]
         for key in rs.keys():
             group.append(rs[key])
         return group

 if __name__ == '__main__':
     u = unionfind([['a', 'b', 'c'], ['b', 'd'], ['e', 'f'], ['g'], ['d', 'h', 'i']])
     u.createtree()
     u.printree()