Statistics
| Branch: | Revision:

iof-tools / networkxMiCe / networkx-master / networkx / utils / union_find.py @ 5cef0f13

History | View | Annotate | Download (3.49 KB)

1
#    Copyright 2016-2019 NetworkX developers.
2
#    Copyright (C) 2004-2019 by
3
#    Aric Hagberg <hagberg@lanl.gov>
4
#    Dan Schult <dschult@colgate.edu>
5
#    Pieter Swart <swart@lanl.gov>
6
#    All rights reserved.
7
#    BSD license.
8
"""
9
Union-find data structure.
10
"""
11

    
12
from networkx.utils import groups
13

    
14

    
15
class UnionFind:
16
    """Union-find data structure.
17

18
    Each unionFind instance X maintains a family of disjoint sets of
19
    hashable objects, supporting the following two methods:
20

21
    - X[item] returns a name for the set containing the given item.
22
      Each set is named by an arbitrarily-chosen one of its members; as
23
      long as the set remains unchanged it will keep the same name. If
24
      the item is not yet part of a set in X, a new singleton set is
25
      created for it.
26

27
    - X.union(item1, item2, ...) merges the sets containing each item
28
      into a single larger set.  If any item is not yet part of a set
29
      in X, it is added to X as one of the members of the merged set.
30

31
      Union-find data structure. Based on Josiah Carlson's code,
32
      http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/215912
33
      with significant additional changes by D. Eppstein.
34
      http://www.ics.uci.edu/~eppstein/PADS/UnionFind.py
35

36
    """
37

    
38
    def __init__(self, elements=None):
39
        """Create a new empty union-find structure.
40

41
        If *elements* is an iterable, this structure will be initialized
42
        with the discrete partition on the given set of elements.
43

44
        """
45
        if elements is None:
46
            elements = ()
47
        self.parents = {}
48
        self.weights = {}
49
        for x in elements:
50
            self.weights[x] = 1
51
            self.parents[x] = x
52

    
53
    def __getitem__(self, object):
54
        """Find and return the name of the set containing the object."""
55

    
56
        # check for previously unknown object
57
        if object not in self.parents:
58
            self.parents[object] = object
59
            self.weights[object] = 1
60
            return object
61

    
62
        # find path of objects leading to the root
63
        path = [object]
64
        root = self.parents[object]
65
        while root != path[-1]:
66
            path.append(root)
67
            root = self.parents[root]
68

    
69
        # compress the path and return
70
        for ancestor in path:
71
            self.parents[ancestor] = root
72
        return root
73

    
74
    def __iter__(self):
75
        """Iterate through all items ever found or unioned by this structure.
76

77
        """
78
        return iter(self.parents)
79

    
80
    def to_sets(self):
81
        """Iterates over the sets stored in this structure.
82

83
        For example::
84

85
            >>> partition = UnionFind('xyz')
86
            >>> sorted(map(sorted, partition.to_sets()))
87
            [['x'], ['y'], ['z']]
88
            >>> partition.union('x', 'y')
89
            >>> sorted(map(sorted, partition.to_sets()))
90
            [['x', 'y'], ['z']]
91

92
        """
93
        # Ensure fully pruned paths
94
        for x in self.parents.keys():
95
            _ = self[x] # Evaluated for side-effect only
96

    
97
        # TODO In Python 3.3+, this should be `yield from ...`.
98
        for block in groups(self.parents).values():
99
            yield block
100

    
101
    def union(self, *objects):
102
        """Find the sets containing the objects and merge them all."""
103
        roots = [self[x] for x in objects]
104
        # Find the heaviest root according to its weight.
105
        heaviest = max(roots, key=lambda r: self.weights[r])
106
        for r in roots:
107
            if r != heaviest:
108
                self.weights[heaviest] += self.weights[r]
109
                self.parents[r] = heaviest