beam: add DSU (union-find) with path compression, rank, size, tests

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

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