Source code for titan_pylib.data_structures.union_find.partial_persistent_union_find

 1from titan_pylib.data_structures.array.partial_persistent_array import (
 2    PartialPersistentArray,
 3)
 4
 5

[docs]
 6class PartialPersistentUnionFind:
 7
 8    def __init__(self, n: int):
 9        self._n: int = n
10        self._parents: PartialPersistentArray[int] = PartialPersistentArray([-1] * n)
11        self._last_time: int = 0
12

[docs]
13    def root(self, x: int, t: int = -1) -> int:
14        assert t == -1 or t <= self._last_time
15        while True:
16            p = self._parents.get(x, t)
17            if p < 0:
18                return x
19            x = p

20

[docs]
21    def unite(self, x: int, y: int, t: int) -> bool:
22        assert t == -1 or t >= self._last_time
23        self._last_time = t
24        x = self.root(x, t)
25        y = self.root(y, t)
26        if x == y:
27            return False
28        if self._parents.get(x, t) > self._parents.get(y, t):
29            x, y = y, x
30        self._parents.set(x, self._parents.get(x, t) + self._parents.get(y, t), t)
31        self._parents.set(y, x, t)
32        return True

33

[docs]
34    def size(self, x: int, t: int = -1) -> int:
35        assert t == -1 or t <= self._last_time
36        return -self._parents.get(self.root(x, t), t)

37

[docs]
38    def same(self, x: int, y: int, t: int = -1) -> bool:
39        assert t == -1 or t <= self._last_time
40        return self.root(x, t) == self.root(y, t)