beam: harden dsu_rollback with range validation and no-compression find

wizard-work/beam/dsu_rollback.py

@@ -26,7 +26,14 @@ class DSU:
         # history: list of (child_root, parent_root, old_size_parent) for undo
         self.history: List[Tuple[int, int, int]] = []
 
+    def _check_index(self, x: int) -> None:
+        if not isinstance(x, int) or isinstance(x, bool):
+            raise TypeError("index must be int")
+        if x < 0 or x >= len(self.parent):
+            raise ValueError("index out of range")
+
     def find(self, x: int) -> int:
+        self._check_index(x)
         # no path compression (required for rollback)
         root = x
         while self.parent[root] != root:
@@ -34,6 +41,8 @@ class DSU:
         return root
 
     def union(self, a: int, b: int) -> bool:
+        self._check_index(a)
+        self._check_index(b)
         ra = self.find(a)
         rb = self.find(b)
         if ra == rb:
@@ -60,9 +69,12 @@ class DSU:
             self.sz[parent] = old_size
 
     def size(self, x: int) -> int:
+        self._check_index(x)
         return self.sz[self.find(x)]
 
     def connected(self, a: int, b: int) -> bool:
+        self._check_index(a)
+        self._check_index(b)
         return self.find(a) == self.find(b)
 
     def components(self) -> int: