ABC420 E — Reachability Query

1# -*- coding: utf-8 -*-2 3 4class UnionFind:5    """Represents a data structure that tracks a set of elements partitioned6       into a number of disjoint (non-overlapping) subsets.7 8    Landau notation: O(α(n)), where α(n) is the inverse Ackermann function.9 10    See:11    https://www.youtube.com/watch?v=zV3Ul2pA2Fw12    https://en.wikipedia.org/wiki/Disjoint-set_data_structure13    https://atcoder.jp/contests/abc120/submissions/444494214    https://atcoder.jp/contests/abc292/submissions/3941007515    https://github.com/not522/ac-library-python/blob/master/atcoder/dsu.py16    """17 18    def __init__(self, number_count: int) -> None:19        """20        Args:21            number_count: The size of elements (greater than 2).22        """23        self.number_count = number_count24        self.parent_numbers = [-1 for _ in range(number_count)]25        self.has_black = [set() for _ in range(number_count)]26 27    def find_root(self, number: int) -> int:28        """Follows the chain of parent pointers from number up the tree until29           it reaches a root element, whose parent is itself.30        Args:31            number: The trees id (0-index).32 33        Returns:34            The index of a root element.35        """36        if self.parent_numbers[number] < 0:37            return number38 39        self.parent_numbers[number] = self.find_root(self.parent_numbers[number])40        return self.parent_numbers[number]41 42    def is_same_group(self, number_x: int, number_y: int) -> bool:43        """Represents the roots of tree number_x and number_y are in the same44           group.45        Args:46            number_x: The trees x (0-index).47            number_y: The trees y (0-index).48        """49        return self.find_root(number_x) == self.find_root(number_y)50 51    def merge_if_needs(self, number_x: int, number_y: int) -> bool:52        """Uses find_root to determine the roots of the tree number_x and53           number_y belong to. If the roots are distinct, the trees are combined54           by attaching the roots of one to the root of the other.55        Args:56            number_x: The trees x (0-index).57            number_y: The trees y (0-index).58        """59        x = self.find_root(number_x)60        y = self.find_root(number_y)61 62        if x == y:63            return False64 65        if self.parent_numbers[x] > self.parent_numbers[y]:66            x, y = y, x67 68        self.parent_numbers[x] += self.parent_numbers[y]69        self.parent_numbers[y] = x70 71        self.has_black[x] |= self.has_black[y]72        return True73 74    def update_color(self, number: int) -> None:75        root = self.find_root(number)76 77        if number in self.has_black[root]:78            self.has_black[root].discard(number)79        else:80            self.has_black[root].add(number)81 82    def is_reachable(self, number: int) -> bool:83        root = self.find_root(number)84 85        if len(self.has_black[root]) >= 1:86            return True87 88        return False89 90 91def main():92    import sys93 94    input = sys.stdin.readline95 96    n, q = map(int, input().split())97    uf = UnionFind(n)98 99    for _ in range(q):100        type, *args = list(map(int, input().split()))101 102        if type == 1:103            u, v = args104            u -= 1105            v -= 1106 107            uf.merge_if_needs(u, v)108        elif type == 2:109            v = args[0] - 1110            uf.update_color(v)111        else:112            v = args[0] - 1113 114            if uf.is_reachable(v):115                print("Yes")116            else:117                print("No")118 119 120if __name__ == "__main__":121    main()122