BZOJ 2588: Spoj 10628. Count on a tree [树上主席树]

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
const int N=1e5+5,INF=1e9+5;
int read(){
    char c=getchar();int x=0,f=1;
    while(c<‘0‘||c>‘9‘){if(c==‘-‘)f=-1; c=getchar();}
    while(c>=‘0‘&&c<=‘9‘){x=x*10+c-‘0‘; c=getchar();}
    return x*f;
}
int n,Q,u,v,k,last,mp[N],m,w[N];
struct ques{
    int u,v,k;
}q[N];
int Bin(int v){
    int l=1,r=m;
    while(l<=r){
        int mid=(l+r)>>1;
        if(mp[mid]==v) return mid;
        else if(v<mp[mid]) r=mid-1;
        else l=mid+1;
    }
    return -1;
}
struct edge{
    int v,ne;
}e[N<<1];
int h[N],cnt;
inline void ins(int u,int v){
    cnt++;
    e[cnt].v=v;e[cnt].ne=h[u];h[u]=cnt;
    cnt++;
    e[cnt].v=u;e[cnt].ne=h[v];h[v]=cnt;
}
int size[N],fa[N],deep[N],mx[N],top[N],tid[N],tot;
void dfs(int u){
    size[u]=1;
    for(int i=h[u];i;i=e[i].ne){
        int v=e[i].v;
        if(v==fa[u]) continue;
        fa[v]=u;deep[v]=deep[u]+1;
        dfs(v);
        size[u]+=size[v];
        if(size[mx[u]]<size[v]) mx[u]=v;
    }
}
void dfs(int u,int anc){
    if(!u) return;
    tid[u]=++tot;
    top[u]=anc;
    dfs(mx[u],anc);
    for(int i=h[u];i;i=e[i].ne){
        int v=e[i].v;
        if(v!=fa[u]&&v!=mx[u]) dfs(v,v);
    }
}
int lca(int x,int y){
    while(top[x]!=top[y]){
        if(deep[top[x]]<deep[top[y]]) swap(x,y);
        x=fa[top[x]];
    }
    if(tid[x]>tid[y]) swap(x,y);
    return x;
}

struct node{
    int l,r,size;
}t[N*20];
int sz,root[N];
void insert(int &x,int l,int r,int num){
    t[++sz]=t[x];x=sz;
    t[x].size++;
    if(l==r) return;
    int mid=(l+r)>>1;
    if(num<=mid) insert(t[x].l,l,mid,num);
    else insert(t[x].r,mid+1,r,num);
}
void build(int u){
    root[u]=root[fa[u]];
    insert(root[u],1,m,Bin(w[u]));
    if(mx[u]) build(mx[u]);
    for(int i=h[u];i;i=e[i].ne)
        if(e[i].v!=fa[u]&&e[i].v!=mx[u]) build(e[i].v);
}
inline int ls(int x){return t[t[x].l].size;}
int query(int u,int v,int k){
    int p=lca(u,v),q=fa[p];
    u=root[u];v=root[v];p=root[p];q=root[q];
    int l=1,r=m;
    while(l!=r){
        int mid=(l+r)>>1,_=ls(u)+ls(v)-ls(p)-ls(q);
        if(k<=_) r=mid,u=t[u].l,v=t[v].l,p=t[p].l,q=t[q].l;
        else l=mid+1,u=t[u].r,v=t[v].r,p=t[p].r,q=t[q].r,k-=_; 
    }
    return l;
}
int main(){
    n=read();Q=read();
    for(int i=1;i<=n;i++) mp[i]=w[i]=read();
    for(int i=1;i<=n-1;i++) u=read(),v=read(),ins(u,v);
    dfs(1);dfs(1,1);
    for(int i=1;i<=Q;i++) q[i].u=read(),q[i].v=read(),q[i].k=read();
    sort(mp+1,mp+1+n);
    m=1;
    for(int i=2;i<=n;i++) if(mp[i]!=mp[i-1]) mp[++m]=mp[i];
    
    build(1);
    for(int i=1;i<=Q;i++){
        u=last^q[i].u;v=q[i].v;k=q[i].k;
        last=mp[query(u,v,k)];
        printf("%d",last);
        if(i!=Q) putchar(‘\n‘);
    }
}