ledger
ca3e13d328
Rebuild test_no_path_compression_find_stable on 4 nodes so the node it asks
about (3) sits two levels down under union-by-size; the old cell built
union(0,1);union(1,2) -> parent [0,0,0,0], so find(2) rewrote nothing and a
path-compressing find passed the cell that names the mechanism.
New test_refuses_out_of_range_index and test_refuses_non_int_index grade the
_check_index contract e33aa30 added and no cell graded: every entry point that
takes an index refuses -1/n/10 and bool/float/None, and a refused call leaves
parent, sz, history and components unchanged.
Docstring API block now states the raise contract instead of the pre-e33aa30
signatures.
Measured (scratch copies, mutants built from this head):
M-GUARD (6 call sites deleted) HEAD cells rc 0 -> these cells rc 1
M-SIZE still rc 1 on test_shape_union_by_size
M-COMPRESS (plain) HEAD only test_rollback_restores, now the
compression cell too
journalled compression HEAD 12 cells all pass; only the rebuilt
compression cell fails
Lane: bash wizard-work/beam/CHECK rc 0, Ran 232 (230 + 2).
88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
"""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)
|