TheAlgorithms_Python/data_structures/disjoint_set/disjoint_set.py at master · AlgorithmAndLeetCode/TheAlgorithms_Python · GitHub

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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""
Disjoint set.
Reference: https://en.wikipedia.org/wiki/Disjoint-set_data_structure
"""

from dataclasses import dataclass


@dataclass
class Node:
    data: int
    rank: int = 0
    parent: Node | None = None


def make_set(x: Node) -> None:
    """
    Make x as a set.

    >>> node = Node(1)
    >>> make_set(node)
    >>> node.parent == node
    True
    >>> node.rank
    0
    >>> node.data
    1
    """
    # rank is the distance from x to its' parent
    # root's rank is 0
    x.rank = 0
    x.parent = x


def union_set(x: Node, y: Node) -> None:
    """
    Union of two sets.
    set with bigger rank should be parent, so that the
    disjoint set tree will be more flat.

    >>> node1 = Node(1)
    >>> node2 = Node(2)
    >>> make_set(node1)
    >>> make_set(node2)
    >>> union_set(node1, node2)
    >>> find_set(node1) == find_set(node2)
    True
    >>> # Test union of already connected nodes
    >>> node3 = Node(3)
    >>> make_set(node3)
    >>> union_set(node1, node3)
    >>> find_set(node1) == find_set(node3)
    True
    >>> find_set(node2) == find_set(node3)
    True
    """
    x, y = find_set(x), find_set(y)
    if x == y:
        return

    elif x.rank > y.rank:
        y.parent = x
    else:
        x.parent = y
        if x.rank == y.rank:
            y.rank += 1


def find_set(x: Node) -> Node:
    """
    Return the parent of x

    >>> node = Node(1)
    >>> make_set(node)
    >>> find_set(node) == node
    True
    >>> node1 = Node(1)
    >>> node2 = Node(2)
    >>> make_set(node1)
    >>> make_set(node2)
    >>> union_set(node1, node2)
    >>> find_set(node1) == find_set(node2)
    True
    >>> # Test path compression
    >>> node3 = Node(3)
    >>> make_set(node3)
    >>> union_set(node1, node3)
    >>> find_set(node1) == find_set(node3)
    True
    """
    if x != x.parent:
        x.parent = find_set(x.parent)
    return x.parent


def find_python_set(node: Node) -> set:
    """
    Return a Python Standard Library set that contains i.

    >>> node = Node(1)
    >>> find_python_set(node)
    {0, 1, 2}
    >>> node = Node(4)
    >>> find_python_set(node)
    {3, 4, 5}
    >>> node = Node(6)
    >>> find_python_set(node)
    Traceback (most recent call last):
       ...
    ValueError: 6 is not in ({0, 1, 2}, {3, 4, 5})
    """
    sets = ({0, 1, 2}, {3, 4, 5})
    for s in sets:
        if node.data in s:
            return s
    msg = f"{node.data} is not in {sets}"
    raise ValueError(msg)


def test_disjoint_set() -> None:
    """
    Test the disjoint set operations with a comprehensive example.
    Creates two disjoint sets: {0, 1, 2} and {3, 4, 5}

    >>> test_disjoint_set()
    """
    vertex = [Node(i) for i in range(6)]
    for v in vertex:
        make_set(v)

    union_set(vertex[0], vertex[1])
    union_set(vertex[1], vertex[2])
    union_set(vertex[3], vertex[4])
    union_set(vertex[3], vertex[5])

    for node0 in vertex:
        for node1 in vertex:
            if find_python_set(node0).isdisjoint(find_python_set(node1)):
                assert find_set(node0) != find_set(node1)
            else:
                assert find_set(node0) == find_set(node1)


if __name__ == "__main__":
    test_disjoint_set()