Spaces:
Running
Running
File size: 3,338 Bytes
b200bda |
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 100 101 102 103 104 105 106 107 |
"""
Union-find data structure.
"""
from networkx.utils import groups
class UnionFind:
"""Union-find data structure.
Each unionFind instance X maintains a family of disjoint sets of
hashable objects, supporting the following two methods:
- X[item] returns a name for the set containing the given item.
Each set is named by an arbitrarily-chosen one of its members; as
long as the set remains unchanged it will keep the same name. If
the item is not yet part of a set in X, a new singleton set is
created for it.
- X.union(item1, item2, ...) merges the sets containing each item
into a single larger set. If any item is not yet part of a set
in X, it is added to X as one of the members of the merged set.
Union-find data structure. Based on Josiah Carlson's code,
https://code.activestate.com/recipes/215912/
with significant additional changes by D. Eppstein.
http://www.ics.uci.edu/~eppstein/PADS/UnionFind.py
"""
def __init__(self, elements=None):
"""Create a new empty union-find structure.
If *elements* is an iterable, this structure will be initialized
with the discrete partition on the given set of elements.
"""
if elements is None:
elements = ()
self.parents = {}
self.weights = {}
for x in elements:
self.weights[x] = 1
self.parents[x] = x
def __getitem__(self, object):
"""Find and return the name of the set containing the object."""
# check for previously unknown object
if object not in self.parents:
self.parents[object] = object
self.weights[object] = 1
return object
# find path of objects leading to the root
path = []
root = self.parents[object]
while root != object:
path.append(object)
object = root
root = self.parents[object]
# compress the path and return
for ancestor in path:
self.parents[ancestor] = root
return root
def __iter__(self):
"""Iterate through all items ever found or unioned by this structure."""
return iter(self.parents)
def to_sets(self):
"""Iterates over the sets stored in this structure.
For example::
>>> partition = UnionFind("xyz")
>>> sorted(map(sorted, partition.to_sets()))
[['x'], ['y'], ['z']]
>>> partition.union("x", "y")
>>> sorted(map(sorted, partition.to_sets()))
[['x', 'y'], ['z']]
"""
# Ensure fully pruned paths
for x in self.parents:
_ = self[x] # Evaluated for side-effect only
yield from groups(self.parents).values()
def union(self, *objects):
"""Find the sets containing the objects and merge them all."""
# Find the heaviest root according to its weight.
roots = iter(
sorted(
{self[x] for x in objects}, key=lambda r: self.weights[r], reverse=True
)
)
try:
root = next(roots)
except StopIteration:
return
for r in roots:
self.weights[root] += self.weights[r]
self.parents[r] = root
|