bzoj2588: Spoj 10628. Count on a tree

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int maxd = 18;
const int maxn = 100000 + 10;
const int maxm = 200000 + 10;
const int maxx = 4000000 + 10;

int n,m,q,la,cnt,size;
int g[maxn],v[maxm],next[maxm],eid;
int anc[maxn][maxd+5];
int h[maxn];
int num[maxn],pos[maxn];
int lc[maxx],rc[maxx],s[maxx];
int a[maxn],t[maxn];
int root[maxn];

void addedge(int a,int b) {
    v[eid]=b; next[eid]=g[a]; g[a]=eid++;
    v[eid]=a; next[eid]=g[b]; g[b]=eid++;    
}

void dfs(int u) {
    num[++cnt]=u; pos[u]=cnt;
    for(int d=1;d<maxd;d++) anc[u][d]=anc[anc[u][d-1]][d-1];
    for(int i=g[u];~i;i=next[i]) if(v[i]!=anc[u][0]) {
        anc[v[i]][0]=u;
        h[v[i]]=h[u]+1;
        dfs(v[i]);    
    }
}

int LCA(int a,int b) {
    if(h[a]<h[b]) swap(a,b);
    for(int i=maxd-1;i>=0;i--) if(h[anc[a][i]]>=h[b]) a=anc[a][i];
    if(a==b) return a;
    for(int i=maxd-1;i>=0;i--) 
        if(anc[a][i]!=anc[b][i]) {
            a=anc[a][i];
            b=anc[b][i];
        }
    return anc[a][0];
}

void modify(int &x,int y,int L,int R,int pos) {
    x=++size; s[x]=s[y]+1;
    lc[x]=lc[y],rc[x]=rc[y];
    if(L==R) return;
    
    int mid=(L+R)>>1;
    if(pos<=mid) modify(lc[x],lc[y],L,mid,pos);
    else modify(rc[x],rc[y],mid+1,R,pos);    
}

int query(int a,int b,int k) {
    int c=LCA(a,b),d=anc[c][0];
    
    a=root[pos[a]]; b=root[pos[b]]; c=root[pos[c]]; d=root[pos[d]];
    int l=1,r=m,mid,cur;
    while(l<r) { 
        mid=(l+r)>>1;
        cur=s[lc[a]]+s[lc[b]]-s[lc[c]]-s[lc[d]];
        if(k<=cur) {
            a=lc[a]; b=lc[b]; c=lc[c]; d=lc[d]; r=mid;    
        }
        else {
            a=rc[a]; b=rc[b]; c=rc[c]; d=rc[d]; l=mid+1; k-=cur;
        }
    }
    return t[l];
}

int main() {
    memset(g,-1,sizeof(g));
    scanf("%d%d",&n,&q);    
    for(int i=1;i<=n;i++) {
        scanf("%d",&a[i]);
        t[i]=a[i];
    }
    sort(t+1,t+n+1);
    m=unique(t+1,t+n+1)-(t+1);
    for(int i=1;i<=n;i++) a[i]=lower_bound(t+1,t+m+1,a[i])-t;
    for(int i=1,u,v;i<n;i++) {
        scanf("%d%d",&u,&v);
        addedge(u,v);
    }
    cnt=0; h[1]=1; anc[1][0]=0;
    dfs(1);
        
    root[0]=lc[0]=rc[0]=s[0]=0;
    for(int i=1;i<=n;i++)
        modify(root[i],root[pos[anc[num[i]][0]]],1,m,a[num[i]]);
    la=0;
    
    for(int i=1,u,v,k;i<=q;i++) {
        scanf("%d%d%d",&u,&v,&k);
        u^=la;
        printf("%d",la=query(u,v,k));
        if(i!=q) printf("\n");
    }
    return 0;
}