class DSU_rollback:
def __init__(self):
self.parent = {}
self.rank = {}
self.history = []
def find(self, v):
if v not in self.parent:
self.parent[v] = v
self.rank[v] = 0
if v == self.parent[v]:
return v
# self.parent[v] = self.find(self.parent[v])
# We don't perform path compression
return self.find(self.parent[v])
def union(self, v, u):
v = self.find(v)
u = self.find(u)
if v == u:
return False
if self.rank[v] > self.rank[u]:
v, u = u, v
self.history.append((v, self.rank[v], u, self.rank[u]))
self.parent[v] = u
if self.rank[v] == self.rank[u]:
self.rank[u] += 1
return True
def rollback(self):
if not self.history:
return
v, rank_v, u, rank_u = self.history.pop()
self.parent[v] = v
self.rank[v] = rank_v
self.parent[u] = u
self.rank[u] = rank_u