diff --git a/wizard-work/beam/dsu.py b/wizard-work/beam/dsu.py new file mode 100644 index 0000000..62538ef --- /dev/null +++ b/wizard-work/beam/dsu.py @@ -0,0 +1,84 @@ +"""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 \ No newline at end of file diff --git a/wizard-work/beam/test_dsu.py b/wizard-work/beam/test_dsu.py new file mode 100644 index 0000000..c9d4162 --- /dev/null +++ b/wizard-work/beam/test_dsu.py @@ -0,0 +1,112 @@ +import random +from dsu import DSU + + +def test_basic(): + d = DSU(5) + assert d.components == 5 + assert not d.connected(0, 1) + assert d.union(0, 1) is True + assert d.connected(0, 1) + assert d.components == 4 + assert d.union(0, 1) is False + assert d.components == 4 + + +def test_size(): + d = DSU(6) + d.union(0, 1); d.union(1, 2) + assert d.size(0) == 3 and d.size(2) == 3 + assert d.size(3) == 1 + d.union(3, 4) + assert d.size(3) == 2 + + +def test_transitive(): + d = DSU(4) + d.union(0, 1); d.union(2, 3); d.union(1, 2) + assert d.components == 1 + for i in range(4): + for j in range(4): + assert d.connected(i, j) + + +def test_empty(): + d = DSU(0) + assert d.components == 0 + assert len(d) == 0 + assert d.sets() == {} + + +def test_reset(): + d = DSU(4) + d.union(0, 1); d.union(2, 3) + d.reset() + assert d.components == 4 + assert not d.connected(0, 1) + + +def test_sets(): + d = DSU(5) + d.union(0, 2); d.union(1, 3) + s = d.sets() + assert sorted(s.values(), key=lambda v: v[0]) == [[0, 2], [1, 3], [4]] + + +def test_errors(): + d = DSU(3) + try: + d.find(5) + except IndexError: + pass + else: + raise AssertionError("expected IndexError") + try: + DSU(-1) + except ValueError: + pass + else: + raise AssertionError("expected ValueError") + try: + DSU(True) + except ValueError: + pass + else: + raise AssertionError("expected ValueError for bool") + + +def test_random_vs_reference(): + rng = random.Random(42) + n = 30 + d = DSU(n) + parent = list(range(n)) + def ref_find(x): + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + for _ in range(2000): + a = rng.randrange(n); b = rng.randrange(n) + # union in both + ra, rb = ref_find(a), ref_find(b) + expected_merge = ra != rb + if expected_merge: + parent[ra] = rb + got = d.union(a, b) + assert got == expected_merge, (a, b, got, expected_merge) + # connectivity check sample + for _ in range(5): + x = rng.randrange(n); y = rng.randrange(n) + assert d.connected(x, y) == (ref_find(x) == ref_find(y)) + assert d.components == len({ref_find(i) for i in range(n)}) + + +def test_deep_chain_no_recursion(): + # union in a line to force deep chains, then find must not recurse + n = 100000 + d = DSU(n) + for i in range(n - 1): + d.union(i, i + 1) + assert d.components == 1 + assert d.find(0) == d.find(n - 1) + assert d.size(0) == n