diff --git a/wizard-work/beam/dsu_rollback.py b/wizard-work/beam/dsu_rollback.py new file mode 100644 index 0000000..e90d1d5 --- /dev/null +++ b/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) diff --git a/wizard-work/beam/test_dsu_rollback.py b/wizard-work/beam/test_dsu_rollback.py new file mode 100644 index 0000000..c9b4f16 --- /dev/null +++ b/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()