Frequency Queries (Codeforces 1591E)
Tags: Others, trees

Petya has a rooted tree with an integer written on each vertex. The vertex $$$1$$$ is the root. You are to answer some questions about the tree.A tree is a connected graph without cycles. A rooted tree has a special vertex called the root. The parent of a node $$$v$$$ is the next vertex on the shortest path from $$$v$$$ to the root.Each question is defined by three integers $$$v$$$, $$$l$$$, and $$$k$$$. To get the answer to the question, you need to perform the following steps:   First, write down the sequence of all integers written on the shortest path from the vertex $$$v$$$ to the root (including those written in the $$$v$$$ and the root).  Count the number of times each integer occurs. Remove all integers with less than $$$l$$$ occurrences.  Replace the sequence, removing all duplicates and ordering the elements by the number of occurrences in the original list in increasing order. In case of a tie, you can choose the order of these elements arbitrary.  The answer to the question is the $$$k$$$-th number in the remaining sequence. Note that the answer is not always uniquely determined, because there could be several orderings. Also, it is possible that the length of the sequence on this step is less than $$$k$$$, in this case the answer is $$$-1$$$. For example, if the sequence of integers on the path from $$$v$$$ to the root is $$$[2, 2, 1, 7, 1, 1, 4, 4, 4, 4]$$$, $$$l = 2$$$ and $$$k = 2$$$, then the answer is $$$1$$$.Please answer all questions about the tree.

Input: Each test contains multiple test cases. The first line contains the number of test cases $$$t$$$ ($$$1 \leq t \leq 10^6$$$). Description of the test cases follows.The first line of each test case contains two integers $$$n$$$, $$$q$$$ ($$$1 \leq n, q \leq 10^6$$$) — the number of vertices in the tree and the number of questions.The second line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$1 \leq a_i \leq n$$$), where $$$a_i$$$ is the number written on the $$$i$$$-th vertex.The third line contains $$$n-1$$$ integers $$$p_2, p_3, \ldots, p_n$$$ ($$$1 \leq p_i \leq n$$$), where $$$p_i$$$ is the parent of node $$$i$$$. It's guaranteed that the values $$$p$$$ define a correct tree.Each of the next $$$q$$$ lines contains three integers $$$v$$$, $$$l$$$, $$$k$$$ ($$$1 \leq v, l, k \leq n$$$) — descriptions of questions.It is guaranteed that the sum of $$$n$$$ and the sum of $$$q$$$ over all test cases do not exceed $$$10^6$$$.

Output: For each question of each test case print the answer to the question. In case of multiple answers, print any.


Sample Inputs:
2
3 3
1 1 1
1 2
3 1 1
3 1 2
3 2 1
5 5
1 2 1 1 2
1 1 2 2
3 1 1
2 1 2
4 1 1
4 2 1
4 2 2

Sample Outputs:
1 -1 1 
1 1 2 1 -1

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m, x):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, 2):
        s[x[i] + 2] += 1
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * m
    for i in range(0, 2 * m, 2):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def f(u, v):
    return u * (n + 1) + v

def fenwick_tree(n):
    tree = [0] * (n + 1)
    return tree

def get_sum(i):
    s = 0
    while i > 0:
        s += tree[i]
        i -= i & -i
    return s

def add(i, x):
    while i < len(tree):
        tree[i] += x
        i += i & -i

def binary_search():
    (x, y) = (0, k)
    for i in range(20, -1, -1):
        if x + pow2[i] < len(tree) and y > tree[x + pow2[i]]:
            x += pow2[i]
            y -= tree[x]
    return x + 1
t = int(input())
ans = []
pow2 = [1]
for _ in range(20):
    pow2.append(2 * pow2[-1])
for _ in range(t):
    (n, q) = map(int, input().split())
    a = [0] + list(map(int, input().split()))
    p = list(map(int, input().split()))
    x = [0] * (2 * (n - 1))
    for i in range(n - 1):
        (x[2 * i], x[2 * i + 1]) = (p[i], i + 2)
    (G, s0) = make_graph(n, n - 1, x)
    (l0, k0) = ([0] * q, [0] * q)
    x = [0] * (2 * q)
    for i in range(q):
        (v, l, k) = map(int, input().split())
        (l0[i], k0[i]) = (l, k)
        (x[2 * i], x[2 * i + 1]) = (v, i)
    (Q, s1) = make_graph(n, q, x)
    ans0 = [-1] * q
    c = [0] * (n + 1)
    c[a[1]] += 1
    tree = fenwick_tree(n + 1)
    add(1, 1)
    s = [0] * (2 * n + 5)
    nx = [-1] * (2 * n + 5)
    p0 = [-1] * (n + 1)
    q0 = [-1] * (n + 1)
    s[0] = a[1]
    (p0[1], q0[1]) = (0, 0)
    st = [1]
    now = list(s0)
    ss = 1
    z = 1
    while st:
        i = st[-1]
        z += 1
        if now[i] == s0[i]:
            for u in range(s1[i], s1[i + 1]):
                (l, k) = (l0[Q[u]], k0[Q[u]])
                ng = get_sum(l - 1)
                k += ng
                if ss >= k:
                    x = binary_search()
                    j = p0[x]
                    while c[s[j]] ^ x:
                        j = nx[j]
                    p0[x] = j
                    ans0[Q[u]] = s[j]
        if now[i] < s0[i + 1]:
            j = G[now[i]]
            st.append(j)
            now[i] += 1
            aj = a[j]
            if c[aj]:
                add(c[aj], -1)
                ss -= 1
            c[aj] += 1
            add(c[aj], 1)
            s[z] = aj
            if p0[c[aj]] == -1:
                p0[c[aj]] = z
            else:
                nx[q0[c[aj]]] = z
            q0[c[aj]] = z
            ss += 1
        else:
            j = st.pop()
            aj = a[j]
            add(c[aj], -1)
            c[aj] -= 1
            ss -= 1
            if c[aj]:
                add(c[aj], 1)
                s[z] = aj
                if p0[c[aj]] == -1:
                    p0[c[aj]] = z
                else:
                    nx[q0[c[aj]]] = z
                q0[c[aj]] = z
                ss += 1
    ans.append(' '.join(map(str, ans0)))
sys.stdout.write('\n'.join(ans))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m, x):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, 2):
        s[x[i] + 2] += 1
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * m
    for i in range(0, 2 * m, 2):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def f(u, v):
    return u * (n + 1) + v

def fenwick_tree(n):
    tree = [0] * (n + 1)
    return tree

def get_sum(i):
    s = 0
    while i > 0:
        s += tree[i]
        i -= i & -i
    return s

def add(i, x):
    while i < len(tree):
        tree[i] += x
        i += i & -i

def binary_search():
    (x, y) = (0, k)
    for i in range(20, -1, -1):
        if x + pow2[i] < len(tree) and y > tree[x + pow2[i]]:
            x += pow2[i]
            y -= tree[x]
    return x + 1
t = int(input())
ans = []
pow2 = [1]
for _ in range(20):
    pow2.append(2 * pow2[-1])
for _ in range(t):
    (n, q) = map(int, input().split())
    a = [0] + list(map(int, input().split()))
    p = list(map(int, input().split()))
    x = [0] * (2 * (n - 1))
    for i in range(n - 1):
        (x[2 * i], x[2 * i + 1]) = (p[i], i + 2)
    (G, s0) = make_graph(n, n - 1, x)
    (l0, k0) = ([0] * q, [0] * q)
    x = [0] * (2 * q)
    for i in range(q):
        (v, l, k) = map(int, input().split())
        (l0[i], k0[i]) = (l, k)
        (x[2 * i], x[2 * i + 1]) = (v, i)
    (Q, s1) = make_graph(n, q, x)
    ans0 = [-1] * q
    c = [0] * (n + 1)
    c[a[1]] += 1
    tree = fenwick_tree(n + 1)
    add(1, 1)
    s = [0] * (2 * n + 5)
    nx = [-1] * (2 * n + 5)
    p0 = [-1] * (n + 1)
    q0 = [-1] * (n + 1)
    s[0] = a[1]
    (p0[1], q0[1]) = (0, 0)
    st = [1]
    now = list(s0)
    ss = 1
    z = 1
    while st:
        i = st[-1]
        z += 1
        if now[i] == s0[i]:
            for u in range(s1[i], s1[i + 1]):
                (l, k) = (l0[Q[u]], k0[Q[u]])
                ng = get_sum(l - 1)
                k += ng
                if ss >= k:
                    x = binary_search()
                    j = p0[x]
                    while c[s[j]] ^ x:
                        j = nx[j]
                    p0[x] = j
                    ans0[Q[u]] = s[j]
        if now[i] < s0[i + 1]:
            j = G[now[i]]
            st.append(j)
            now[i] += 1
            aj = a[j]
            if c[aj]:
                add(c[aj], -1)
                ss -= 1
            c[aj] += 1
            add(c[aj], 1)
            s[z] = aj
            if p0[c[aj]] == -1:
                p0[c[aj]] = z
            else:
                nx[q0[c[aj]]] = z
            q0[c[aj]] = z
            ss += 1
        else:
            j = st.pop()
            aj = a[j]
            add(c[aj], -1)
            c[aj] -= 1
            ss -= 1
            if c[aj]:
                add(c[aj], 1)
                s[z] = aj
                if p0[c[aj]] == -1:
                    p0[c[aj]] = z
                else:
                    nx[q0[c[aj]]] = z
                q0[c[aj]] = z
                ss += 1
    ans.append(' '.join(map(str, ans0)))
sys.stdout.write('\n'.join(ans))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m, x):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, 2):
        s[x[i] + 2] += 1
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * m
    for i in range(0, 2 * m, 2):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def f(u, v):
    return u * (n + 1) + v

def fenwick_tree(n):
    tree = [0] * (n + 1)
    return tree

def get_sum(i):
    s = 0
    while i > 0:
        s += tree[i]
        i -= i & -i
    return s

def add(i, x):
    while i < len(tree):
        tree[i] += x
        i += i & -i

def binary_search():
    (x, y) = (0, k)
    for i in range(20, -1, -1):
        if x + pow2[i] < len(tree) and y > tree[x + pow2[i]]:
            x += pow2[i]
            y -= tree[x]
    return x + 1
t = int(input())
ans = []
pow2 = [1]
for _ in range(20):
    pow2.append(2 * pow2[-1])
for _ in range(t):
    (n, q) = map(int, input().split())
    a = [0] + list(map(int, input().split()))
    p = list(map(int, input().split()))
    x = [0] * (2 * (n - 1))
    for i in range(n - 1):
        (x[2 * i], x[2 * i + 1]) = (p[i], i + 2)
    (G, s0) = make_graph(n, n - 1, x)
    vlk = [tuple(map(int, input().split())) for _ in range(q)]
    x = [0] * (2 * q)
    for i in range(q):
        (v, _, _) = vlk[i]
        (x[2 * i], x[2 * i + 1]) = (v, i)
    (Q, s1) = make_graph(n, q, x)
    ans0 = [-1] * q
    c = [0] * (n + 1)
    c[a[1]] += 1
    tree = fenwick_tree(n + 1)
    add(1, 1)
    s = [0] * (2 * n + 5)
    nx = [-1] * (2 * n + 5)
    p0 = [-1] * (n + 1)
    q0 = [-1] * (n + 1)
    s[0] = a[1]
    (p0[1], q0[1]) = (0, 0)
    st = [1]
    now = list(s0)
    ss = 1
    z = 1
    while st:
        i = st[-1]
        z += 1
        if now[i] == s0[i]:
            for u in range(s1[i], s1[i + 1]):
                (v, l, k) = vlk[Q[u]]
                ng = get_sum(l - 1)
                k += ng
                if ss >= k:
                    x = binary_search()
                    j = p0[x]
                    while c[s[j]] ^ x and j < nx[j]:
                        j = nx[j]
                    p0[x] = j
                    ans0[Q[u]] = s[j]
        if now[i] < s0[i + 1]:
            j = G[now[i]]
            st.append(j)
            now[i] += 1
            aj = a[j]
            if c[aj]:
                add(c[aj], -1)
                ss -= 1
            c[aj] += 1
            add(c[aj], 1)
            s[z] = aj
            if p0[c[aj]] == -1:
                p0[c[aj]] = z
            else:
                nx[q0[c[aj]]] = z
            q0[c[aj]] = z
            ss += 1
        else:
            j = st.pop()
            aj = a[j]
            add(c[aj], -1)
            c[aj] -= 1
            ss -= 1
            if c[aj]:
                add(c[aj], 1)
                s[z] = aj
                if p0[c[aj]] == -1:
                    p0[c[aj]] = z
                else:
                    nx[q0[c[aj]]] = z
                q0[c[aj]] = z
                ss += 1
    ans.append(' '.join(map(str, ans0)))
sys.stdout.write('\n'.join(ans))