disjointset.py

• 이전 버전 코드

from typing import TypeVar

T = TypeVar('T')

# N DisjointSet
# I {"version": "2.0"}
class DisjointSet:
    """Disjoint Set for integers with path compression and union-by-size."""

    def __init__(self, size: int):
        self._parent = [-1] * size

    def union(self, x: int, y: int, should_raise: bool = False) -> int:
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            if should_raise:
                raise ValueError(f'{x} and {y} are already in the same set.')
            else:
                return root_x
        if self._parent[root_x] > self._parent[root_y]:
            root_x, root_y = root_y, root_x
        self._parent[root_x] += self._parent[root_y]
        self._parent[root_y] = root_x
        return root_x

    def find(self, x: int) -> int:
        while (p := self._parent[x]) >= 0:
            x, self._parent[x] = p, self._parent[p]
        return x

• Disjoint Set 참고