Union-Find Tree
(byslib/data/union_find.py)

• View this file on GitHub
• Last update: 2022-04-18 22:29:22+09:00

Verified with

• tests/data/union_find.test.py
• tests/graph/kruskal.test.py

Code

# @brief Union-Find Tree
from typing import Dict, List


class UnionFindTree:
    """Union-Find Tree
    Notes
    -----
    Time complexity

    * union : :math:`O(α(N))`
    * find : :math:`O(α(N))`

    Examples
    --------
    >>> uft = UnionFindTree(5)
    >>> uft.union(0, 4)
    True
    >>> uft.union(0, 3)
    True
    >>> print(uft.same(3, 4))
    True

    References
    ----------
    ..[1] 🐜 p.81
    """

    def __init__(self, n: int) -> None:
        self.n = n
        self.n_tree = n
        self.parent = [-1] * self.n  # 負なら親でありその値は(-子の数)

    def find(self, a: int) -> int:
        now = a
        path = []
        while self.parent[now] >= 0:
            path.append(now)
            now = self.parent[now]
        else:
            root = now

        for p in path:
            self.parent[p] = root

        return root

    def union(self, a: int, b: int) -> bool:
        root_a = self.find(a)
        root_b = self.find(b)

        if root_a == root_b:
            return False

        if -self.parent[root_a] > -self.parent[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] += self.parent[root_a]
        self.parent[root_a] = root_b

        self.n_tree -= 1
        return True

    def same(self, a: int, b: int) -> bool:
        return self.find(a) == self.find(b)

    def size(self, a: int) -> int:
        return -self.parent[self.find(a)]

    def group(self) -> Dict[int, List[int]]:
        res: Dict[int, List[int]] = dict()
        for i in range(self.n):
            leader = self.find(i)
            if leader in res:
                res[leader].append(i)
            else:
                res[leader] = [i]

        return res


if __name__ == "__main__":
    import doctest

    doctest.testmod()

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.10.4/x64/lib/python3.10/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir], 'release': True}).decode()
  File "/opt/hostedtoolcache/Python/3.10.4/x64/lib/python3.10/site-packages/onlinejudge_verify/languages/python.py", line 80, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page