Source code for titan_pylib.data_structures.union_find.persistent_union_find

 1from titan_pylib.data_structures.array.persistent_array import PersistentArray
 2from typing import Optional
 3
 4

[docs]
 5class PersistentUnionFind:
 6
 7    def __init__(self, n: int, _parents: Optional[PersistentArray[int]] = None) -> None:
 8        """``n`` 個の要素からなる ``PersistentUnionFind`` を構築します。
 9        :math:`O(n)` です。
10        """
11        self._n: int = n
12        self._parents: PersistentArray[int] = (
13            PersistentArray([-1] * n) if _parents is None else _parents
14        )
15
16    def _new(self, _parents: PersistentArray[int]) -> "PersistentUnionFind":
17        return PersistentUnionFind(self._n, _parents)
18

[docs]
19    def copy(self) -> "PersistentUnionFind":
20        """コピーします。
21        :math:`O(1)` です。
22        """
23        return self._new(self._parents.copy())

24

[docs]
25    def root(self, x: int) -> int:
26        """要素 ``x`` を含む集合の代表元を返します。
27        :math:`O(\\log^2{n})` です。
28        """
29        _parents = self._parents
30        while True:
31            p = _parents.get(x)
32            if p < 0:
33                return x
34            x = p

35

[docs]
36    def unite(self, x: int, y: int, update: bool = True) -> "PersistentUnionFind":
37        """要素 ``x`` を含む集合と要素 ``y`` を含む集合を併合します。
38        :math:`O(\\log^2{n})` です。
39
40        Args:
41          x (int): 集合の要素です。
42          y (int): 集合の要素です。
43          update (bool, optional): 併合後を新しいインスタンスにするなら ``True`` です。
44
45        Returns:
46          PersistentUnionFind: 併合後の uf です。
47        """
48        x = self.root(x)
49        y = self.root(y)
50        res_parents = self._parents.copy() if update else self._parents
51        if x == y:
52            return self._new(res_parents)
53        px, py = res_parents.get(x), res_parents.get(y)
54        if px > py:
55            x, y = y, x
56        res_parents = res_parents.set(x, px + py)
57        res_parents = res_parents.set(y, x)
58        return self._new(res_parents)

59

[docs]
60    def size(self, x: int) -> int:
61        """要素 ``x`` を含む集合の要素数を返します。
62        :math:`O(\\log^2{n})` です。
63        """
64        return -self._parents.get(self.root(x))

65

[docs]
66    def same(self, x: int, y: int) -> bool:
67        """
68        要素 ``x`` と ``y`` が同じ集合に属するなら ``True`` を、
69        そうでないなら ``False`` を返します。
70        :math:`O(\\log^2{n})` です。
71        """
72        return self.root(x) == self.root(y)