ABC372 E — K-th Largest Connected Components

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.members = [[i] for i 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        self.members[x] += self.members[y]71        self.members[x] = sorted(self.members[x], reverse=True)[:10]72        return True73 74 75def main():76    import sys77 78    input = sys.stdin.readline79 80    n, q = map(int, input().split())81    uf = UnionFind(n)82 83    for _ in range(q):84        qi = list(map(int, input().split()))85 86        if qi[0] == 1:87            u, v = qi[1:]88            u -= 189            v -= 190 91            if not uf.is_same_group(u, v):92                uf.merge_if_needs(u, v)93        else:94            vi, k = qi[1:]95            vi -= 196 97            root_vi = uf.find_root(vi)98            members = uf.members[root_vi]99 100            if len(members) >= k:101                print(members[k - 1] + 1)102            else:103                print(-1)104 105 106if __name__ == "__main__":107    main()108