"""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) n < 0 -> ValueError find(x) -> root no path compression; TypeError if x is not an int or is a bool, ValueError if x is outside [0, n) union(a, b) -> bool True if merged; both arguments are checked as above and a refused call mutates nothing snapshot() -> int (marker) rollback(marker) -> None revert to marker; TypeError if marker is not an int or is a bool, ValueError if marker outside [0, len(history)] size(x) -> component size index checked as find connected(a, b) -> bool both indexes checked as find """ 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 _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: root = self.parent[root] 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: # 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 isinstance(marker, bool) or not isinstance(marker, int): raise TypeError("marker must be int") 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: 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: # number of distinct roots return sum(1 for i in range(len(self.parent)) if self.parent[i] == i)