Standard Union-Find doesn't support "undo." For problems needing rollback, maintain a history stack:
class RollbackUnionFind:
def __init__(self, n):
self.parent = list(range(n))
self.rank = [0] * n
self.history = [] # (node, old_parent, old_rank)
def union(self, x, y):
rootX, rootY = self.find(x), self.find(y)
if rootX == rootY:
return False
# Store state before union
self.history.append((rootX, self.parent[rootX], self.rank[rootX]))
self.history.append((rootY, self.parent[rootY], self.rank[rootY]))
# Perform union
...
def rollback(self):
node, parent, rank = self.history.pop()
self.parent[node] = parent
self.rank[node] = rank
Important: rollback Union-Find cannot use path compression (it makes history tracking complex). You get operations.