ロールバック付きUnion-Find

union by size だけを行い経路圧縮をしないことで、merge で書き換えた値を履歴に積んでおき、あとから取り消せるようにしたものです。

• RollbackUnionFind(N): $N$ 頂点で初期化
• merge(a, b): $a$ と $b$ をつなぐ。すでに同じ集合なら False（このときも履歴には積むので、undo 1 回で取り消せる）
• undo(): 最後の merge を取り消す。 $O(1)$
• snapshot(): 今の状態を表す値を返す。 $O(1)$
• rollback(t): snapshot() が $t$ を返したときの状態に戻す（既定は $t = 0$ で、何もつないでいない状態）。取り消す merge 1 回につき $O(1)$
• leader(a), same(a, b), size(a): 普通の Union-Find と同じ。 $O(\log N)$

辺の追加と削除が混ざるクエリを、時間を区間に分けたセグ木の上で DFS して処理する（オフライン動的連結性）ときなどに使います。

例:

uf = RollbackUnionFind(n)
t = uf.snapshot()
for a, b in edges:
    uf.merge(a, b)
# ここで edges を加えたときの連結性を調べる
uf.rollback(t)  # edges を加える前に戻る

コード

class RollbackUnionFind:
    def __init__(self, n):
        self.p = [-1] * n
        self.history = []  # merge のたびに (つないだ側の根, その元の p) を積む

    def leader(self, a):
        while self.p[a] >= 0:
            a = self.p[a]
        return a

    def merge(self, a, b):
        a, b = self.leader(a), self.leader(b)
        if a == b:
            self.history.append((-1, 0))
            return False
        if self.p[a] > self.p[b]:
            a, b = b, a
        self.history.append((b, self.p[b]))
        self.p[a] += self.p[b]
        self.p[b] = a
        return True

    def undo(self):
        # 最後の merge を取り消す
        b, pb = self.history.pop()
        if b == -1:
            return
        a = self.p[b]
        self.p[b] = pb
        self.p[a] -= pb

    def snapshot(self):
        # 今の状態を表す値（rollback に渡す）
        return len(self.history)

    def rollback(self, t=0):
        # snapshot() が t を返したときの状態に戻す
        while len(self.history) > t:
            self.undo()

    def same(self, a, b):
        return self.leader(a) == self.leader(b)

    def size(self, a):
        return -self.p[self.leader(a)]