Przejdź do treści

Struktura zbiorów rozłącznych

Opis problemu

Implementacja

 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
class DisjointUnion:
    """
    Class representing Disjoint-Set-Union structure.
    :param number_of_nodes: number of elements stored in the structure.
    :ivar _subsets: information about subsets.
    """

    class Node:
        def __init__(self, parent: int, rank: int):
            self.parent = parent
            self.rank = rank

    def __init__(self, number_of_nodes: int):
        self._subsets = [self.Node(i, 0) for i in range(number_of_nodes)]

    def union(self, el1: int, el2: int) -> None:
        """
        Union two elements.
        :param el1: First element.
        :param el2: Second element.
        :return: None
        """
        x_root = self.find(el1)
        y_root = self.find(el2)

        if self._subsets[x_root].rank < self._subsets[y_root].rank:
            self._subsets[x_root].parent = y_root
        elif self._subsets[x_root].rank > self._subsets[y_root].rank:
            self._subsets[y_root].parent = x_root
        else:
            self._subsets[y_root].parent = x_root
            self._subsets[x_root].rank += 1

    def is_in_union(self, el1: int, el2: int) -> bool:
        """
        Checks if two elements are in union i.e. have the same representative.
        :param el1: First element.
        :param el2: Second element.
        :return: True if x and y are in union, False otherwise.
        """
        return self.find(el1) == self.find(el2)

    def find(self, node_number: int):
        """
        Find the representative of a given element.
        :param node_number: element to check
        :return: Representative of the node_number element in the structure.
        """
        if self._subsets[node_number].parent != node_number:
            self._subsets[node_number].parent = self.find(self._subsets[node_number].parent)

        return self._subsets[node_number].parent

if __name__ == "__main__":
    disjoint_union = DisjointUnion(10)

    disjoint_union.union(0, 1)

    print(disjoint_union.is_in_union(0, 1))
    print(disjoint_union.is_in_union(1, 2))