84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
"""Disjoint Set Union (union-find) with path compression and union by rank.
|
|
|
|
Provides amortized near-constant find/union, component counting, and
|
|
size queries. Iterative find avoids Python recursion limits on deep chains.
|
|
|
|
Public API:
|
|
DSU(n) -> n singleton sets (indices 0..n-1)
|
|
find(x) -> representative of x's set
|
|
union(a, b) -> True if merged (were distinct), else False
|
|
connected(a, b) -> bool
|
|
size(x) -> size of x's component
|
|
components -> number of disjoint sets
|
|
reset() -> restore to singletons
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
class DSU:
|
|
__slots__ = ("_parent", "_rank", "_size", "_components")
|
|
|
|
def __init__(self, n: int):
|
|
if not isinstance(n, int) or isinstance(n, bool) or n < 0:
|
|
raise ValueError("n must be a non-negative int")
|
|
self._parent = list(range(n))
|
|
self._rank = [0] * n
|
|
self._size = [1] * n
|
|
self._components = n
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._parent)
|
|
|
|
@property
|
|
def components(self) -> int:
|
|
return self._components
|
|
|
|
def find(self, x: int) -> int:
|
|
if x < 0 or x >= len(self._parent):
|
|
raise IndexError("index out of range")
|
|
# iterative path compression (two-pass)
|
|
root = x
|
|
while self._parent[root] != root:
|
|
root = self._parent[root]
|
|
# compress path
|
|
while self._parent[x] != root:
|
|
nxt = self._parent[x]
|
|
self._parent[x] = root
|
|
x = nxt
|
|
return root
|
|
|
|
def union(self, a: int, b: int) -> bool:
|
|
ra = self.find(a)
|
|
rb = self.find(b)
|
|
if ra == rb:
|
|
return False
|
|
if self._rank[ra] < self._rank[rb]:
|
|
ra, rb = rb, ra
|
|
self._parent[rb] = ra
|
|
self._size[ra] += self._size[rb]
|
|
if self._rank[ra] == self._rank[rb]:
|
|
self._rank[ra] += 1
|
|
self._components -= 1
|
|
return True
|
|
|
|
def connected(self, a: int, b: int) -> bool:
|
|
return self.find(a) == self.find(b)
|
|
|
|
def size(self, x: int) -> int:
|
|
return self._size[self.find(x)]
|
|
|
|
def reset(self) -> None:
|
|
n = len(self._parent)
|
|
self._parent = list(range(n))
|
|
self._rank = [0] * n
|
|
self._size = [1] * n
|
|
self._components = n
|
|
|
|
def sets(self):
|
|
"""Return dict rep -> sorted list of members (O(n alpha(n)))."""
|
|
groups = {}
|
|
for i in range(len(self._parent)):
|
|
r = self.find(i)
|
|
groups.setdefault(r, []).append(i)
|
|
return groups |