1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
|
# This is a general algorithm used by the annotator, translator, and other code
# union-find impl, a info object is attached to the roots
class UnionFind(object):
def __init__(self, info_factory=None):
self.link_to_parent = {}
self.weight = {}
self.info_factory = info_factory
self.root_info = {}
# mapping-like [] access
def __getitem__(self, obj):
if obj not in self.link_to_parent:
raise KeyError(obj)
ignore, rep, info = self.find(obj)
return info
def __contains__(self, obj):
return obj in self.link_to_parent
def __iter__(self):
return iter(self.link_to_parent)
def keys(self):
return self.link_to_parent.keys()
def infos(self):
return self.root_info.values()
def find_rep(self, obj):
try:
# fast path (shortcut for performance reasons)
parent = self.link_to_parent[obj]
self.root_info[parent] # may raise KeyError
return parent
except KeyError:
# general case
ignore, rep, info = self.find(obj)
return rep
def find(self, obj): # -> new_root, obj, info
if obj not in self.link_to_parent:
if self.info_factory:
info = self.info_factory(obj)
else:
info = None
self.root_info[obj] = info
self.weight[obj] = 1
self.link_to_parent[obj] = obj
return True, obj, info
to_root = [obj]
parent = self.link_to_parent[obj]
while parent is not to_root[-1]:
to_root.append(parent)
parent = self.link_to_parent[parent]
for obj in to_root:
self.link_to_parent[obj] = parent
return False, parent, self.root_info[parent]
def union(self, obj1, obj2): # -> not_noop, rep, info
new1, rep1, info1 = self.find(obj1)
new2, rep2, info2 = self.find(obj2)
if rep1 is rep2:
return new1 or new2, rep1, info1
if info1 is not None:
info1.absorb(info2)
w1 = self.weight[rep1]
w2 = self.weight[rep2]
w = w1 + w2
if w1 < w2:
rep1, rep2 = rep2, rep1
self.link_to_parent[rep2] = rep1
del self.weight[rep2]
del self.root_info[rep2]
self.weight[rep1] = w
self.root_info[rep1] = info1
return True, rep1, info1
def union_list(self, objlist):
if len(objlist) == 0:
return
obj0 = objlist[0]
self.find(obj0)
for obj1 in objlist[1:]:
self.union(obj0, obj1)
|