Union-Find木

class UnionFind:
    def __init__(self, n):
        self.table = [-1] * n

    def root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x

    def is_same(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
        else:
            self.table[r1] = r2
            self.table[r2] += d1

    def get_size(self, x):
        return -self.table[self.root(x)]

class UnionFindWithGrouping:
    def __init__(self, n):
        self.table = [-1] * n
        self.group = [{i} for i in range(n)]

    def root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x

    def is_same(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            return

        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
            self.group[r1].update(self.group[r2])
        else:
            self.table[r1] = r2
            self.table[r2] += d1
            self.group[r2].update(self.group[r1])

    def get_size(self, x):
        return -self.table[self.root(x)]

    def get_group(self, x):
        return self.group[self.root(x)]

class RollbackUnionFind:
    """
    履歴を持ち、直前の操作を（順に）取り消せるUnionFind
    """
 
    def __init__(self, n: int):
        self.table = [-1] * n
        self.history = []
        # (継続リーダー番号1, 被統合リーダー番号2, 1統合前table値, 2統合前table値)
        # 無効なuniteの場合、被統合リーダー番号を"-1"とする
 
    def root(self, x: int):
        tbl = self.table
        while tbl[x] >= 0:
            x = tbl[x]
        return x
 
    def is_same(self, x: int, y: int):
        return self.root(x) == self.root(y)
 
    def unite(self, x: int, y: int):
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            self.history.append((r1, -1, -1, -1))
            return False
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
            self.history.append((r1, r2, d1, d2))
        else:
            self.table[r1] = r2
            self.table[r2] += d1
            self.history.append((r2, r1, d2, d1))
        return True
 
    def rollback(self):
        if not self.history:
            return False
        r1, r2, d1, d2 = self.history.pop()
        if r2 == -1:
            return False
        self.table[r1] = d1
        self.table[r2] = d2
        return True
 
    def get_size(self, x: int):
        return -self.table[self.root(x)]

部分永続②の方。未検証。

from bisect import bisect


class PartialPersistentUnionFind:
    """
    時刻 t は、「t番目のuniteクエリ完了時点」を意味する。初期状態は t=0。

    root(x [, t]): 時刻 t の x の根（t 省略時は最新）
    find(x, y [, t]): 時刻 t に x と y が同一グループか判定
    unite(x, y): x と y を結合
    get_size(x [, t]): 時刻 t の x の属するグループの要素数
    """

    def __init__(self, n: int):
        self.t = 0
        self.timetable = [[0] for _ in range(n)]
        self.valuetable = [[-1] for _ in range(n)]

    def _getvalue(self, x: int, t: int):
        i = bisect(self.timetable[x], t) - 1
        return self.valuetable[x][i]

    def root(self, x: int, t: int = -1):
        # 更新が行われるのは、その時点の根に対してだけなので、
        # x の最終更新時刻 > t なら、t 時点では確実に x は根。
        # そうでない場合は最終更新時の値の正負で、通常のUnionFindと同様、根か、親を持っているか判別。
        if t == -1:
            t = self.t
        while self.timetable[x][-1] <= t and self.valuetable[x][-1] >= 0:
            x = self.valuetable[x][-1]
        return x

    def is_same(self, x: int, y: int, t: int = -1):
        return self.root(x, t) == self.root(y, t)

    def unite(self, x: int, y: int):
        self.t += 1
        r1 = self.root(x, self.t)
        r2 = self.root(y, self.t)
        if r1 == r2:
            return False
        self.timetable[r1].append(self.t)
        self.timetable[r2].append(self.t)
        d1 = self.valuetable[r1][-1]
        d2 = self.valuetable[r2][-1]
        if d1 <= d2:
            self.valuetable[r1].append(d1 + d2)
            self.valuetable[r2].append(r1)
        else:
            self.valuetable[r1].append(r2)
            self.valuetable[r2].append(d1 + d2)
        return True

    def get_size(self, x: int, t: int = -1):
        if t == -1:
            t = self.t
        r = self.root(x, t)
        return -self._getvalue(r, t)

from typing import TypeVar, Generic, List, Optional, Sequence

T = TypeVar('T')


class PersistentArrayNode(Generic[T]):
    def __init__(self, m: int, t: int, value: Optional[T] = None):
        self.children: List[Optional['PersistentArrayNode']] = [None] * m
        self.created: int = t
        self.value: Optional[T] = value

    def copy(self, t: int):
        res = PersistentArrayNode(0, t, self.value)
        res.children = self.children.copy()  # shallow copy
        return res


class PersistentArray(Generic[T]):
    """
    永続配列
    """

    def __init__(self, array: Sequence[T]):
        n = len(array)
        log_n = (n - 1).bit_length()
        root = PersistentArrayNode(log_n, 0, array[0])

        self.n = n
        self.m = log_n  # M分木を構築
        self.roots = [root]

        # 初期化
        q = [(0, 1, root)]  # (そのノードが示すi, 隣り合う子の差分, ノードインスタンス)
        m = self.m
        while q:
            i, d, node = q.pop()
            nd = d * m
            # children[0] だけはleading-zero問題があるので、別処理
            # （現在の i=xxxx(M進数) としたとき、0xxxx に進むので、そのノード自身の値を持たない）
            ci = i + nd
            if ci < n:
                child = PersistentArrayNode(log_n, 0)
                node.children[0] = child
                q.append((i, nd, child))
            for j in range(1, m):
                ci = i + d * j
                if ci < n:
                    child = PersistentArrayNode(log_n, 0, array[ci])
                    node.children[j] = child
                    q.append((ci, nd, child))
                else:
                    break

    def get_node(self, i: int, t: int) -> PersistentArrayNode:
        # a[i] を示すノードを取得
        # https://37zigen.com/persistent-array/
        # （データの持ち方: 2番目の方法, M進数にして下位の桁から見る）
        node = self.roots[t]
        m = self.m
        while i > 0:
            node = node.children[i % m]
            i //= m
        return node

    def get_value(self, i: int, t: int) -> T:
        return self.get_node(i, t).value

    def update(self, i: int, x: T, t: int) -> int:
        t_new = len(self.roots)
        node = self.roots[t].copy(t_new)
        self.roots.append(node)
        m = self.m

        while i > 0:
            child = node.children[i % m].copy(t_new)
            node.children[i % m] = child
            node = child
            i //= m

        node.value = x

        return t_new

    def update_newest(self, i: int, x: T):
        """
        最新バージョンに対し、新規バージョンを作らないで更新する。
        複数箇所を一括に更新することがあり、一括更新の途中状態は不要という場合、
        最初だけ update し、2番目以降の更新をこちらにすることで、tの進み方が呼び出し元とずれることを防ぐ。
        """
        # ノードに自身か作成された時刻を持たせ、作成時刻が最新バージョンより古ければコピー、同じなら使い回す
        t_new = len(self.roots) - 1
        node = self.roots[-1]
        m = self.m

        while i > 0:
            child = node.children[i % m]
            if child.created < t_new:
                child = child.copy(t_new)
                node.children[i % m] = child
            node = child
            i //= m

        node.value = x

    def debug_print(self):
        for t in range(len(self.roots)):
            res = []
            for i in range(self.n):
                res.append(self.get_value(i, t))
            print(t, res)


class PersistentUnionFind(Generic[T]):

    def __init__(self, n):
        self.table = PersistentArray([-1] * n)

    def root(self, x, t):
        tbl = self.table
        p = tbl.get_value(x, t)
        while p >= 0:
            x = p
            p = tbl.get_value(x, t)
        return x

    def find(self, x, y, t):
        return self.root(x, t) == self.root(y, t)

    def unite(self, x, y, t):
        r1 = self.root(x, t)
        r2 = self.root(y, t)
        if r1 == r2:
            return False
        d1 = self.table.get_value(r1, t)
        d2 = self.table.get_value(r2, t)
        if d1 <= d2:
            self.table.update(r2, r1, t)
            self.table.update_newest(r1, d1 + d2)
        else:
            self.table.update(r1, r2, t)
            self.table.update_newest(r2, d1 + d2)
        return True

    def get_size(self, x, t):
        return -self.table.get_value(self.root(x, t), t)

from typing import Callable, Generic, TypeVar, List

T = TypeVar('T')


class UnionFindWithPotential(Generic[T]):

    def __init__(self,
                 n: int,
                 init: Callable[[], T],
                 func: Callable[[T, T], T],
                 rev_func: Callable[[T, T], T]):
        """
        :param n:
        :param init: 単位元の生成関数
        :param func: 2項間加算関数（add）
        :param rev_func: 逆関数（sub）
        """
        self.table: List[int] = [-1] * n
        self.values: List[T] = [init() for _ in range(n)]
        self.init: Callable[[], T] = init
        self.func: Callable[[T, T], T] = func
        self.rev_func: Callable[[T, T], T] = rev_func

    def root(self, x: int) -> int:
        stack = []
        tbl = self.table
        vals = self.values

        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        if stack:
            val = self.init()
            while stack:
                y = stack.pop()
                val = self.func(val, vals[y])
                vals[y] = val
                tbl[y] = x
        return x

    def is_same(self, x: int, y: int) -> bool:
        return self.root(x) == self.root(y)

    def diff(self, x: int, y: int) -> T:
        """
        x と y の差（y - x）を取得。同じグループに属さない場合は None。
        """
        if not self.is_same(x, y):
            return None
        vx = self.values[x]
        vy = self.values[y]
        return self.rev_func(vy, vx)

    def unite(self, x: int, y: int, d: T) -> bool:
        """
        x と y のグループを、y - x = d となるように統合。
        既に x と y が同グループで、矛盾する場合は AssertionError。矛盾しない場合はFalse。
        同グループで無く、新たな統合が発生した場合はTrue。
        """
        rx = self.root(x)
        ry = self.root(y)
        vx = self.values[x]
        vy = self.values[y]
        if rx == ry:
            assert self.rev_func(vy, vx) == d
            return False

        rd = self.rev_func(self.func(vx, d), vy)
        self.table[rx] += self.table[ry]
        self.table[ry] = rx
        self.values[ry] = rd
        return True

    def get_size(self, x: int) -> int:
        return -self.table[self.root(x)]