beam: add DSU with rollback (union-by-size, no path compression)

wizard-work/beam/dsu_rollback.py

@@ -0,0 +1,70 @@
+"""Disjoint Set Union with rollback (persistent, undoable).
+
+Subtle point: rollback requires union by SIZE/RANK only — path compression
+is incompatible with undo because it mutates parents in a way that cannot be
+cheaply reversed. This implementation uses union by size and records enough
+state to undo each union in O(1).
+
+API:
+  DSU(n)
+  find(x) -> root (no path compression)
+  union(a, b) -> bool  (True if merged)
+  snapshot() -> int (marker)
+  rollback(marker) -> None (revert to marker)
+  size(x) -> component size
+  connected(a, b) -> bool
+"""
+from typing import List, Tuple
+
+
+class DSU:
+    def __init__(self, n: int):
+        if n < 0:
+            raise ValueError("n must be non-negative")
+        self.parent: List[int] = list(range(n))
+        self.sz: List[int] = [1] * n
+        # history: list of (child_root, parent_root, old_size_parent) for undo
+        self.history: List[Tuple[int, int, int]] = []
+
+    def find(self, x: int) -> int:
+        # no path compression (required for rollback)
+        root = x
+        while self.parent[root] != root:
+            root = self.parent[root]
+        return root
+
+    def union(self, a: int, b: int) -> bool:
+        ra = self.find(a)
+        rb = self.find(b)
+        if ra == rb:
+            # record a no-op marker? No — nothing changed; nothing to undo.
+            return False
+        # union by size: attach smaller to larger
+        if self.sz[ra] < self.sz[rb]:
+            ra, rb = rb, ra
+        # now ra is larger (or equal); attach rb under ra
+        self.history.append((rb, ra, self.sz[ra]))
+        self.parent[rb] = ra
+        self.sz[ra] += self.sz[rb]
+        return True
+
+    def snapshot(self) -> int:
+        return len(self.history)
+
+    def rollback(self, marker: int) -> None:
+        if marker < 0 or marker > len(self.history):
+            raise ValueError("invalid marker")
+        while len(self.history) > marker:
+            child, parent, old_size = self.history.pop()
+            self.parent[child] = child
+            self.sz[parent] = old_size
+
+    def size(self, x: int) -> int:
+        return self.sz[self.find(x)]
+
+    def connected(self, a: int, b: int) -> bool:
+        return self.find(a) == self.find(b)
+
+    def components(self) -> int:
+        # number of distinct roots
+        return sum(1 for i in range(len(self.parent)) if self.parent[i] == i)

wizard-work/beam/test_dsu_rollback.py

@@ -0,0 +1,95 @@
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import unittest
+from dsu_rollback import DSU
+
+
+class TestDSURollback(unittest.TestCase):
+    def test_basic_union_connected(self):
+        d = DSU(4)
+        self.assertTrue(d.union(0, 1))
+        self.assertTrue(d.connected(0, 1))
+        self.assertFalse(d.connected(0, 2))
+        self.assertEqual(d.size(0), 2)
+
+    def test_noop_union_returns_false_no_history(self):
+        d = DSU(3)
+        d.union(0, 1)
+        m = d.snapshot()
+        self.assertFalse(d.union(0, 1))
+        self.assertEqual(d.snapshot(), m)
+
+    def test_rollback_restores(self):
+        d = DSU(5)
+        d.union(0, 1); d.union(2, 3)
+        m = d.snapshot()
+        d.union(1, 2); d.union(3, 4)
+        self.assertTrue(d.connected(0, 4))
+        d.rollback(m)
+        self.assertFalse(d.connected(0, 4))
+        self.assertTrue(d.connected(0, 1))
+        self.assertTrue(d.connected(2, 3))
+        self.assertFalse(d.connected(1, 2))
+
+    def test_nested_snapshots(self):
+        d = DSU(6)
+        d.union(0, 1)
+        m1 = d.snapshot()
+        d.union(2, 3)
+        m2 = d.snapshot()
+        d.union(1, 2)
+        d.rollback(m2)
+        self.assertTrue(d.connected(0, 1) and d.connected(2, 3))
+        self.assertFalse(d.connected(0, 2))
+        d.rollback(m1)
+        self.assertTrue(d.connected(0, 1))
+        self.assertFalse(d.connected(2, 3))
+
+    def test_full_rollback_to_zero(self):
+        d = DSU(4)
+        d.union(0, 1); d.union(2, 3); d.union(1, 2)
+        d.rollback(0)
+        for i in range(4):
+            self.assertEqual(d.size(i), 1)
+        self.assertEqual(d.components(), 4)
+
+    def test_invalid_marker(self):
+        d = DSU(2)
+        with self.assertRaises(ValueError):
+            d.rollback(5)
+        with self.assertRaises(ValueError):
+            d.rollback(-1)
+
+    def test_no_path_compression_find_stable(self):
+        d = DSU(4)
+        d.union(0, 1); d.union(1, 2)
+        before = list(d.parent)
+        r = d.find(2)
+        self.assertEqual(r, d.find(0))
+        self.assertEqual(d.parent[2], before[2])
+
+    def test_size_after_rollback(self):
+        d = DSU(3)
+        d.union(0, 1)
+        m = d.snapshot()
+        d.union(0, 2)
+        self.assertEqual(d.size(0), 3)
+        d.rollback(m)
+        self.assertEqual(d.size(0), 2)
+        self.assertEqual(d.size(2), 1)
+
+    def test_components(self):
+        d = DSU(4)
+        self.assertEqual(d.components(), 4)
+        d.union(0, 1); d.union(2, 3)
+        self.assertEqual(d.components(), 2)
+        d.union(1, 2)
+        self.assertEqual(d.components(), 1)
+
+    def test_negative_n(self):
+        with self.assertRaises(ValueError):
+            DSU(-1)
+
+
+if __name__ == "__main__":
+    unittest.main()