Count on a tree SPOJ - COT

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<stack>
#include<cstdlib>
#include<queue>
#include<set>
#include<string.h>
#include<vector>
#include<deque>
#include<map>
using namespace std;
#define INF 0x3f3f3f3f3f3f3f3f
#define inf 0x3f3f3f3f
#define eps 1e-4
#define bug printf("*********\n")
#define debug(x) cout<<#x"=["<<x<<"]" <<endl
typedef long long LL;
typedef long long ll;
const int maxn = 2e5 + 5;
const int mod = 998244353;

int n,q,m,TOT;
int a[maxn],t[maxn],T[maxn],lson[maxn*30],rson[maxn*30],c[maxn*30];
void Init_hash() {
    for (int i = 1; i <= n; i++) t[i] = a[i];
    sort(t + 1, t + 1 + n);
    m = unique(t + 1, t + 1 + n) - t - 1;
}
int build(int l,int r) {
    int root = TOT++;
    c[root] = 0;
    if (l != r) {
        int mid = (l + r) >> 1;
        lson[root] = build(l, mid);
        rson[root] = build(mid + 1, r);
    }
    return root;
}
int Hash(int x) {
    return lower_bound(t + 1, t + 1 + m, x) - t;
}
int update(int root,int pos,int val)
{
    int newroot = TOT ++,tmp = newroot;
    c[newroot] = c[root] + val;
    int l = 1,r = m;
    while(l <r)
    {
        int mid = (l+r)>>1;
        if(pos <= mid)
        {
            lson[newroot] = TOT++;
            rson[newroot] = rson[root];
            newroot = lson[newroot];
            root = lson[root];
            r = mid;
        }
        else
        {
            rson[newroot] = TOT ++;
            lson[newroot] = lson[root];
            newroot = rson[newroot];
            root = rson[root];
            l = mid + 1;
        }
        c[newroot] = c[root] + val;
    }
    return tmp;

}
int query(int left_root,int right_root,int LCA,int k) {
    int lca_root = T[LCA];
    int pos = Hash(a[LCA]);
    int l = 1, r = m;
    while (l < r) {
        int mid = (l + r) >> 1;
        int tmp = c[lson[left_root]] + c[lson[right_root]] - 2 * c[lson[lca_root]] + (pos >= l && pos <= mid);
        if (tmp >= k) {
            left_root = lson[left_root];
            right_root = lson[right_root];
            lca_root = lson[lca_root];
            r = mid;
        } else {
            k -= tmp;
            left_root = rson[left_root];
            right_root = rson[right_root];
            lca_root = rson[lca_root];
            l = mid + 1;
        }
    }
    return l;
}
//*********************LAC**************************
int rmq[2 * maxn];  // 欧拉序列对应的深度序列
struct ST {
    int mm[2 * maxn];
    int dp[2 * maxn][20];

    void init(int n) {
        mm[0] = -1;
        for (int i = 1; i <= n; i++) {
            mm[i] = ((i & (i - 1)) == 0) ? mm[i - 1] + 1 : mm[i - 1];
            dp[i][0] = i;
        }
        for (int j = 1; j <= mm[n]; j++)
            for (int i = 1; i + (1 << j) - 1 <= n; i++)
                dp[i][j] =
                        rmq[dp[i][j - 1]] < rmq[dp[i + (1 << (j - 1))][j - 1]] ? dp[i][j - 1] : dp[i + (1 << (j - 1))][
                                j - 1];
    }

    int query(int a, int b) {
        if (a > b) swap(a, b);
        int k = mm[b - a + 1];
        return rmq[dp[a][k]] <= rmq[dp[b - (1 << k) + 1][k]] ? dp[a][k] : dp[b - (1 << k) + 1][k];
    }
}st;
struct Edge {
    int to, next;
}edge[maxn * 2];
int tot,head[maxn * 2];

int F[maxn * 2];    //欧拉序列 即dfs遍历的顺序
int P[maxn];        //表示点i在F中第一次出现的位置
int cnt;
void init() {
    tot = 0;
    memset(head, -1, sizeof head);
}
void addedge(int u,int v) {
    edge[tot].to = v;
    edge[tot].next = head[u];
    head[u] = tot++;
}
void dfs(int u,int pre,int dep) {
    F[++cnt] = u;
    rmq[cnt] = dep;
    P[u] = cnt;
    for (int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].to;
        if (v == pre) continue;
        dfs(v, u, dep + 1);
        F[++cnt] = u;
        rmq[cnt] = dep;
    }
}
void LCA_init(int root,int node_num) {
    cnt = 0;
    dfs(root, root, 0);
    st.init(2 * node_num - 1);
}
int query_lca(int u,int v) {
    return F[st.query(P[u], P[v])];
}

void dfs_build(int u,int pre) {
    int pos = Hash(a[u]);
    T[u] =update(T[pre],pos,1);
    for(int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].to;
        if(v == pre) continue;
        dfs_build(v,u);
    }
}
int main() {
    while(scanf("%d %d",&n,&q) == 2) {
        for(int i = 1; i <= n; i++) scanf("%d",&a[i]);
        Init_hash();
        init();
        TOT = 0;
        int u,v;
        for(int i = 1;i < n; i ++) {
            scanf("%d %d",&u,&v);
            addedge(u,v);
            addedge(v,u);
        }
        LCA_init(1,n);
        T[n + 1] = build(1,m);
        dfs_build(1,n + 1);
        int k;
        while(q--) {
            scanf("%d %d %d",&u,&v,&k);
            int ans = t[query(T[u],T[v],query_lca(u,v),k)];
            printf("%d\n",ans);
        }
        return 0;
    }
    return 0;
}