Approach summary:- Use union-by-rank and path compression, but record every parent/rank/size change on a stack so we can undo (rollback) to a previous state. Each union pushes a marker and the changed entries; rollback pops until marker and restores old values.- Note trade-off: path compression gives near-constant find but may record many changes; for heavy rollback use, prefer union-by-rank without path compression (guaranteed O(log n) finds). Below implementation supports both (path compression by default, can disable).python
class RollbackDSU:
def __init__(self, n, compress=True):
self.n = n
self.parent = list(range(n))
self.rank = [0]*n
self.size = [1]*n
self.changes = [] # stack of tuples (attr, index, old_value)
self.markers = [] # positions for union operations
self.compress = compress
def find(self, x):
if self.parent[x] == x:
return x
if not self.compress:
# no path compression: easier to rollback, O(log n) with union-by-rank
return self.find(self.parent[x])
# path compression with logging
root = self.find(self.parent[x])
if self.parent[x] != root:
self.changes.append(('parent', x, self.parent[x]))
self.parent[x] = root
return root
def union(self, a, b):
a = self.find(a); b = self.find(b)
# marker to group changes for this union
self.markers.append(len(self.changes))
if a == b:
# record a no-op marker by pushing a special marker length unchanged
return False
# union by rank
if self.rank[a] < self.rank[b]:
a, b = b, a
# attach b under a; record changes
self.changes.append(('parent', b, self.parent[b]))
self.parent[b] = a
self.changes.append(('size', a, self.size[a]))
self.size[a] += self.size[b]
if self.rank[a] == self.rank[b]:
self.changes.append(('rank', a, self.rank[a]))
self.rank[a] += 1
return True
def snapshot(self):
# return current change-stack length as a handle for rollback
return len(self.changes)
def rollback_to(self, snapshot_len):
# undo until we reach snapshot_len
while len(self.changes) > snapshot_len:
typ, idx, old = self.changes.pop()
if typ == 'parent':
self.parent[idx] = old
elif typ == 'rank':
self.rank[idx] = old
elif typ == 'size':
self.size[idx] = old
def component_size(self, x):
return self.size[self.find(x)]
Key points and reasoning:- Every mutation (parent, rank, size) is logged so rollback restores exact previous state.- snapshot() returns a point to rollback to; union pushes changes; rollback_to restores.Time complexity guarantees:- With path compression + union by rank: amortized inverse-Ackermann α(n) per find/union (very close to O(1)). However, because we log every parent change, a single find with path compression may produce O(depth) logged changes — rollback cost grows accordingly.- If compress=False (no path compression) and union-by-rank: find is O(log n) worst-case, union O(1), rollback O(1) per change (each union makes O(1) changes). This is the common choice for offline dynamic connectivity because it bounds amount of work to rollback.- Rollback_to cost is linear in number of logged changes undone; if you rollback per union, amortized cost is O(α(n)) with compression or O(1) without compression.Edge cases:- Repeated unions of same component are no-ops (we don't log changes).- Ensure snapshot/rollback boundaries are managed by caller (e.g., divide-and-conquer offline dynamic connectivity).Trade-offs:- Use compress=True for fastest queries if rollbacks are rare.- Use compress=False for many rollbacks (predictable cost and simpler guarantees).