Persistent Segment Tree - COT Count on a tree

/**
  * Persistent Segment Tree + Binary Search + Fast IO = O(n logn logn) - TLE
  * Persistent Segment Tree + Query + Fast IO         = O(n logn) - AC
  * Note            : Works for duplicate values
**/
#include <bits/stdc++.h>

using namespace std;

static const int MAXN = 1e5 + 5;
static const int LOGN = 19;

map <int, int> valTOind;
int indTOval[MAXN];

int tNode, tQuery, n;
int nodeVal[MAXN];// temp[MAXN];
vector <int> graph[MAXN];

int readInt()
{
    bool minus = false;
    int result = 0;
    char ch;
    ch = getchar();
    while (true)
    {
        if (ch == '-')
            break;
        if (ch >= '0' && ch <= '9') break;
            ch = getchar();
    }
    if (ch == '-')
        minus = true; else result = ch-'0';
    while (true)
    {
        ch = getchar();
        if (ch < '0' || ch > '9')
            break;
        result = result*10 + (ch - '0');
    }
    if (minus)
        return -result;
    else
        return result;
}

struct node
{
    int val;
    node *left, *right;
    node(int val)
    {
        this->val = val;
        this->left = NULL;
        this->right = NULL;
    }
} *version[MAXN<<1];

void build(node *root, int a, int b)
{
    if (a > b)
        return;
    if (a == b)
    {
        root->val = 0;
        return;
    }
    int mid = (a+b)>>1;
    root->left = new node(0);
    root->right = new node(0);
    build(root->left, a, mid);
    build(root->right, mid+1, b);
    root->val = root->left->val + root->right->val;
}

void update(node *proot, node *root, int a, int b, int pos)
{
    if (a > b || a > pos || b < pos)
        return;
    if (a >= pos && b <= pos)
    {
        root->val += 1;
        return;
    }
    int mid = (a+b)>>1;
    if (pos <= mid)
    {
        root->left = new node(0);
        root->right = proot->right;
        update(proot->left, root->left, a, mid, pos);
    }
    else
    {
        root->left = proot->left;
        root->right = new node(0);
        update(proot->right, root->right, mid+1, b, pos);
    }
    root->val = root->left->val + root->right->val;
}
/**
int query(node *root, int a, int b, int i, int j)
{
    if (a > b || a > j || b < i)
        return 0;
    if (a >= i && b <= j)
        return root->val;
    int mid = (a+b)>>1;
    int p1 = query(root->left, a, mid, i, j);
    int p2 = query(root->right, mid+1, b, i, j);
    return p1+p2;
}
int Count(node *root_u, node *root_v, node *root_lca, node *root_plca, int k)
{
    int sum = 0;
        sum += query(root_u, 1, n, 1, k);
        sum += query(root_v, 1, n, 1, k);
        sum -= query(root_lca, 1, n, 1, k);
        sum -= query(root_plca, 1, n, 1, k);
    return sum;
}
int binarySearch(node *root_u, node *root_v, node *root_lca, node *root_plca,
                 int k)  // l : version[u], r : version[v]
{
    int low = 1;
    int high = n;
    int ans;
    while (low <= high)
    {
        int mid = (low+high)>>1;
        int cnt = Count(root_u, root_v, root_lca, root_plca, mid);
        if (cnt >= k)
        {
            ans = mid;
            high = mid-1;
        }

        else
        {
            low = mid+1;
        }
    }
    return ans;
}
**/
int query1(node *root_u, node *root_v, node *root_lca, node *root_plca, int a,
           int b, int k)
{
    if (a == b)
        return a;

    int sum = 0;
        sum += root_u->left->val;
        sum += root_v->left->val;
        sum -= root_lca->left->val;
        sum -= root_plca->left->val;
    int mid = (a+b)>>1;
    if (sum >= k)
        return query1(root_u->left, root_v->left, root_lca->left, root_plca->left,
                      a, mid, k);
    else
        return query1(root_u->right, root_v->right, root_lca->right,
                      root_plca->right, mid+1, b, k - sum);
}
int father[MAXN][LOGN];
int depth[MAXN];

void dfs(int u, int p = -1)
{
    for (int i = 1; i < LOGN; i++)
        father[u][i] = father[father[u][i-1]][i-1];
    for (int v : graph[u])
    {
        if (v == p)
            continue;
        father[v][0] = u;
        depth[v] = depth[u] + 1;
        dfs(v, u);
    }
}

int LCA(int u, int v)
{
    if (depth[u] < depth[v])
        swap(u, v);
    for (int i = LOGN-1; i >= 0; i--)
    {
        if (depth[father[u][i]] >= depth[v])
        {
            u = father[u][i];
        }
    }
    if (u == v)
        return u;
    for (int i = LOGN-1; i >= 0; i--)
    {
        if (father[u][i] != father[v][i])
        {
            u = father[u][i];
            v = father[v][i];
        }
    }
    return father[u][0];
}

int ver[MAXN];
bool vis[MAXN];

void bfs(int src)
{
    memset(vis, 0, sizeof vis);
    queue <int> PQ;
    PQ.push(src);
    int vr = 1;
    ver[src] = vr;
    vis[src] = 1;
    version[ver[src]] = new node(0);
    update(version[0], version[ver[src]], 1, n, nodeVal[src]);
    while (!PQ.empty())
    {
        int u = PQ.front(); PQ.pop();
        for (int v : graph[u])
        {
            if (vis[v])
                continue;
            vis[v] = 1;
            vr++;
            ver[v] = vr;
            version[ver[v]] = new node(0);
            update(version[ver[u]], version[ver[v]], 1, n, nodeVal[v]);
            PQ.push(v);
        }
    }

}

struct structure
{
    int val, ind;
    structure() {}
    structure(int val, int ind)
    {
        this->val = val;
        this->ind = ind;
    }
    friend bool operator<(structure A, structure B)
    {
        if (A.val == B.val)
            return A.ind < B.ind;
        else
            return A.val < B.val;
    }
} temp[MAXN];

int main()
{
    tNode = readInt();
    tQuery = readInt();

    for (int i = 1; i <= tNode; i++)
    {
        nodeVal[i] = readInt();
        temp[i] = {nodeVal[i], i};
    }
    for (int e = 1; e < tNode; e++)
    {
        int u, v;
        u = readInt();
        v = readInt();
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    sort(temp+1, temp+tNode+1);
    int id = 0;
    for (int i = 1; i <= tNode; i++)
    {
        int num = temp[i].val;
        int ind = temp[i].ind;
        id++;
        nodeVal[ind] = id;
        indTOval[id] = num;
    }
    n = id;
    version[0] = new node(0);
    build(version[0], 1, n);
    bfs(1);
    depth[1] = 1;
    dfs(1);
    while (tQuery--)
    {
        int u, v, k;
        u = readInt();
        v = readInt();
        k = readInt();
        int lca = LCA(u, v);
        int plca = father[lca][0];
        //int ansb = binarySearch(version[ver[u]], version[ver[v]],
                                  version[ver[lca]], version[ver[plca]], k);
        int ans = query1(version[ver[u]], version[ver[v]], version[ver[lca]],
                         version[ver[plca]], 1, n, k);
        printf("%d\n", indTOval[ans]);
    }
}