CCO Preparation Test 1 P3 - K-th Rank Student.py

# MLE, CHECK C++ CODE WHICH USE PBDS INSTEAD OF SORTEDLIST
#
# https://dmoj.ca/problem/ccoprep1p3
# Store each group with a sorted structure, so we can efficiently get the k-th smallest element
# When joining groups, we use a disjoint set to keep track of each groups' root, also use
# union by size (merge smaller groups to larger ones) so fewer elements are moved (moving an element takes log(n))
#
# TC: O(N*log^2(N) + Qlog(N))

from bisect import bisect_left as lower_bound
from bisect import bisect_right as upper_bound
import sys

input = sys.stdin.readline
print = lambda x: sys.stdout.write(str(x) + "\n")


class FenwickTree:
    def __init__(self, x):
        bit = self.bit = list(x)
        size = self.size = len(bit)
        for i in range(size):
            j = i | (i + 1)
            if j < size:
                bit[j] += bit[i]

    def update(self, idx, x):
        """updates bit[idx] += x"""
        while idx < self.size:
            self.bit[idx] += x
            idx |= idx + 1

    def __call__(self, end):
        """calc sum(bit[:end])"""
        x = 0
        while end:
            x += self.bit[end - 1]
            end &= end - 1
        return x

    def find_kth(self, k):
        """Find largest idx such that sum(bit[:idx]) <= k"""
        idx = -1
        for d in reversed(range(self.size.bit_length())):
            right_idx = idx + (1 << d)
            if right_idx < self.size and self.bit[right_idx] <= k:
                idx = right_idx
                k -= self.bit[idx]
        return idx + 1, k


class SortedList:
    block_size = 700

    def __init__(self, iterable=()):
        self.macro = []
        self.micros = [[]]
        self.micro_size = [0]
        self.fenwick = FenwickTree([0])
        self.size = 0
        for item in iterable:
            self.insert(item)

    def insert(self, x):
        i = lower_bound(self.macro, x)
        j = upper_bound(self.micros[i], x)
        self.micros[i].insert(j, x)
        self.size += 1
        self.micro_size[i] += 1
        self.fenwick.update(i, 1)
        if len(self.micros[i]) >= self.block_size:
            self.micros[i:i + 1] = self.micros[i][:self.block_size >> 1], self.micros[i][self.block_size >> 1:]
            self.micro_size[i:i + 1] = self.block_size >> 1, self.block_size >> 1
            self.fenwick = FenwickTree(self.micro_size)
            self.macro.insert(i, self.micros[i + 1][0])

    def pop(self, k=-1):
        i, j = self._find_kth(k)
        self.size -= 1
        self.micro_size[i] -= 1
        self.fenwick.update(i, -1)
        return self.micros[i].pop(j)

    def __getitem__(self, k):
        i, j = self._find_kth(k)
        return self.micros[i][j]

    def count(self, x):
        return self.upper_bound(x) - self.lower_bound(x)

    def __contains__(self, x):
        return self.count(x) > 0

    def lower_bound(self, x):
        i = lower_bound(self.macro, x)
        return self.fenwick(i) + lower_bound(self.micros[i], x)

    def upper_bound(self, x):
        i = upper_bound(self.macro, x)
        return self.fenwick(i) + upper_bound(self.micros[i], x)

    def _find_kth(self, k):
        return self.fenwick.find_kth(k + self.size if k < 0 else k)

    def __len__(self):
        return self.size

    def __iter__(self):
        return (x for micro in self.micros for x in micro)

    def __repr__(self):
        return str(list(self))


class DisjointSet:
    def __init__(self, N, arr):
        self.parent = [i for i in range(N + 1)]  # disjoint set stuff
        self.size = [1] * (N + 1)
        self.groups =[SortedList([i]) for i in arr]

    def find(self, node):
        if self.parent[node] != node:  # go up until we reach the root
            # attach everything on our way to the root (path compression)
            self.parent[node] = self.find(self.parent[node])
        return self.parent[node]

    def union(self, a, b):
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a == root_b:
            return
        if self.size[root_b] > self.size[root_a]:  # join the smaller to the larger
            root_a, root_b = root_b, root_a  # swap
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]
        for i in self.groups[root_b]:
            self.groups[root_a].insert(i)


N, M = map(int, input().split())
rank = [-1] + list(map(int, input().split()))

loc = [-1]*(N+1)
for i,v in enumerate(rank):
    loc[v] = i

ds = DisjointSet(N, rank)
for _ in range(M):
    a, b = map(int, input().split())
    ds.union(a, b)

Q = int(input())
for _ in range(Q):
    q = input().split()
    if q[0] == "B":
        x, y = map(int, q[1:])
        ds.union(x, y)

    else:
        x, k = map(int, q[1:])
        root = ds.find(x)
        if len(ds.groups[root]) < k:
            print(-1)
        else:
            # print(ds.groups[root])
            print(loc[ds.groups[root][k - 1]])