poj 2104 C - Count on a tree

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<set>
#include<string.h>

#define MAXN  1000100
#define  ll long long
using namespace std;
#define Mod 998244353

int rt[20*MAXN],ls[20*MAXN],rs[20*MAXN],sum[20*MAXN];
int z[MAXN],b[MAXN];
int sz,cnt;

void Build(int l,int r,int &a)
{
    a=++cnt;
    sum[a]=0;
    if(l==r)
        return ;
    int mid=(l+r)>>1;
    Build(l,mid,ls[a]);
    Build(mid+1,r,rs[a]);
}

void Update(int l,int r,int last,int ind,int &a)
{
    a=++cnt;
    ls[a]=ls[last];
    rs[a]=rs[last];
    sum[a]=sum[last]+1;
    if(l==r)
        return ;
    int mid=(l+r)>>1;
    if(ind<=mid)
        Update(l,mid,ls[a],ind,ls[a]);
    else
        Update(mid+1,r,rs[a],ind,rs[a]);
}

int Ques(int l,int r,int l1,int r1,int k)
{
    if(l==r)
        return l;
    int s=sum[ls[r1]]-sum[ls[l1]];
    int mid=(l+r)>>1;
    if(k<=s)
        return Ques(l,mid,ls[l1],ls[r1],k);
    else
        return Ques(mid+1,r,rs[l1],rs[r1],k-s);
}

int main()
{
    int t;
    scanf("%d",&t);
    while(t--)
    {
        int n,q;
        scanf("%d%d",&n,&q);
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&z[i]);
            b[i]=z[i];
        }
        sort(b+1,b+n+1);
        sz=unique(b+1,b+n+1)-(b+1);
        cnt=0;
        Build(1,sz,rt[0]);
        for(int i=1;i<=n;i++)
            z[i]=lower_bound(b+1,b+sz+1,z[i])-b;//这边是sz 
        for(int i=1;i<=n;i++)
            Update(1,sz,rt[i-1],z[i],rt[i]);
        while(q--)
        {
            int l,r,k;
            scanf("%d%d%d",&l,&r,&k);
            int ans=Ques(1,sz,rt[l-1],rt[r],k);
            printf("%d\n",b[ans]);
        }
    }
    return 0;
}

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<set>
#include<string.h>

#define MAXN  200100
#define  ll long long
using namespace std;
#define Mod 998244353

int z[MAXN],b[MAXN],de[MAXN],head[MAXN],rt[MAXN*20],ls[MAXN*20],rs[MAXN*20],sum[MAXN*20],anc[MAXN][25],f[MAXN];
struct node
{
    int v,next;
}edge[MAXN<<1];
int cnt,tot,sz;

void add(int u,int v)
{
    edge[cnt].v=v;
    edge[cnt].next=head[u];
    head[u]=cnt++;
}


int lca(int a,int b)
{
    if(de[a]<de[b])
        swap(a,b);
    for(int i=22;i>=0;i--)
        if(de[b]<=de[anc[a][i]])
            a=anc[a][i];
    if(a==b)
        return a;
    for(int i=22;i>=0;i--)
        if(anc[a][i]!=anc[b][i])
        {
            a=anc[a][i];
            b=anc[b][i];
        }
    return anc[a][0];
}



void Build(int l,int r,int &a)
{
    a=++tot;
    sum[a]=0;
    int mid=(l+r)>>1;
    if(l==r)
        return ;
    Build(l,mid,ls[a]);
    Build(mid+1,r,rs[a]);
}

void Update(int l,int r,int last,int ind,int &a)
{
    a=++tot;
    ls[a]=ls[last];
    rs[a]=rs[last];
    sum[a]=sum[last]+1;
    if(l==r)
        return ;
    int mid=(l+r)>>1;
    if(ind<=mid)
        Update(l,mid,ls[a],ind,ls[a]);
    else
        Update(mid+1,r,rs[a],ind,rs[a]);
}
void dfs(int u,int fa)
{

    de[u]=de[fa]+1;
    anc[u][0]=fa;
    f[u]=fa;
 // cout<<u<<" "<<fa<<endl;
    Update(1,sz,rt[fa],z[u],rt[u]);
    for(int i=1;i<=22;i++)
        anc[u][i]=anc[anc[u][i-1]][i-1];
    
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==fa)
            continue;
        dfs(v,u);
    }
}

int Ques(int l,int r,int l1,int r1,int lc,int  flc,int k)
{
    if(l==r)
        return l;
    int s=sum[ls[l1]]+sum[ls[r1]]-sum[ls[lc]]-sum[ls[flc]];
    int mid=(l+r)>>1;
    if(k<=s)
        return Ques(l,mid,ls[l1],ls[r1],ls[lc],ls[flc],k);
    else
        return Ques(mid+1,r,rs[l1],rs[r1],rs[lc],rs[flc],k-s);

}

int main()
{
    int n,m;

    while(scanf("%d%d",&n,&m)!=EOF)
    {
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&z[i]);
            b[i]=z[i];
        }
        memset(de,0,sizeof(de));
        memset(head,-1,sizeof(head));
        memset(anc,0,sizeof(anc));
        sort(b+1,b+n+1);
        sz=unique(b+1,b+n+1)-(b+1);
        cnt=tot=0;
        Build(1,sz,rt[0]);

        for(int i=1;i<n;i++)
        {
            int a,b;
            scanf("%d%d",&a,&b);
            add(a,b);
            add(b,a);
        }
        for(int i=1;i<=n;i++)
            z[i]=lower_bound(b+1,b+cnt+1,z[i])-b;
        dfs(1,0);
        while(m--)
        {
            int a,c,k;
            scanf("%d%d%d",&a,&c,&k);
            int d=lca(a,c);
            int ind=Ques(1,sz,rt[a],rt[c],rt[d],rt[f[d]],k);
            printf("%d\n",b[ind]);
        }
    }

    return 0;
}