SP1487 PT07J - Query on a tree III 主席树+dfs序

Code:

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<string>
#include<cstring>
#define REP(i,a,n)for(int i=a;i<=n;++i)
#define CLR(d,a)memset(d,a,sizeof(d));

using namespace std;

void SetIO(string a){
    string in=a+".in";
    freopen(in.c_str(),"r",stdin);
}

const int maxn=100000+5;

int val[maxn], Sorted[maxn], n,edges;

int idx[maxn];

void Discrete(){
	REP(i,1,n) Sorted[i]=val[i];
	sort(Sorted+1,Sorted+1+n);
	REP(i,1,n){
		val[i]=lower_bound(Sorted+1,Sorted+1+n,val[i])-Sorted;
		idx[val[i]]=i;
	}
}

int head[maxn<<1],to[maxn<<1],nex[maxn<<1];

void add_edge(int u,int v){
	nex[++edges]=head[u];
	head[u]=edges;
	to[edges]=v;
}

void Read(){
	scanf("%d",&n);
	REP(i,1,n) scanf("%d",&val[i]);
	REP(i,1,n-1){
		int a,b;
		scanf("%d%d",&a,&b);
		add_edge(a,b);
		add_edge(b,a);
	}
}

int arr[maxn], siz[maxn],nodes,position[maxn];

int dfs(int u,int fa){
	siz[u]=1;
	arr[++nodes]=u;
	position[u]=nodes;

	for(int v=head[u];v;v=nex[v]){
		if(to[v]==fa)continue;
		siz[u]+=dfs(to[v],u);
	}
	return siz[u];
}

const int const_Tree=70;

int numv[maxn*const_Tree],root[maxn];

struct Chair_Tree{
	int cnt_Tree,lson[maxn*const_Tree],rson[maxn*const_Tree];

	void build(int l,int r,int &o){
		if(l>r)return;
		o=++cnt_Tree;
		if(l==r)return;
		int mid=(l+r)>>1;
		build(l,mid,lson[o]);
		build(mid+1,r,rson[o]);
	}

	int insert(int l,int r,int o,int pos){
		int oo=++cnt_Tree;
		numv[oo]=numv[o]+1;
		lson[oo]=lson[o];
		rson[oo]=rson[o];
		if(l==r)return oo;

		int mid=(l+r)>>1;
		if(pos<=mid)
			lson[oo]=insert(l,mid,lson[o],pos);
		else
			rson[oo]=insert(mid+1,r,rson[o],pos);
		return oo;
	}

	int query(int l,int r,int pre,int cur,int k){
		if(l==r)return l;
		int lsum=numv[lson[cur]]-numv[lson[pre]];
		int mid=(l+r)>>1;

		if(k<=lsum)
			return query(l,mid,lson[pre],lson[cur],k);
		else
			return query(mid+1,r,rson[pre],rson[cur],k-lsum);
	}
}Tree;

void Build(){
	Discrete();
	dfs(1,0);
	Tree.build(1,n,root[0]);

	REP(i,1,n){
		int u=arr[i];
		root[i]=Tree.insert(1,n,root[i-1],val[u]);
	}
}

void Work(){
	int m;
	scanf("%d",&m);
	REP(i,1,m){
		int a,k;
		scanf("%d%d",&a,&k);
		int root1=root[position[a]-1];
		int root2=root[position[a]+siz[a]-1];

		int ans=Tree.query(1,n,root1,root2,k);
		ans=idx[ans];
		printf("%d\n",ans);
	}
}

int main(){
	SetIO("input");
	Read();
	Build();
	Work();
	return 0;
}