BZOJ_1146_[CTSC2008]网络管理Network_主席树+树状数组

#include <stdio.h>
#include <string.h>
#include <algorithm>
using namespace std;
#define N 80050
#define maxn 100000000
int n,m,head[N],to[N<<1],nxt[N<<1],val[N],dfn[N],cnt,out[N];
int fa[N],top[N],son[N],size[N],dep[N],tot,A[N],B[N],C[N],D[N];
int siz[N*250],ls[N*250],rs[N*250],root[N];
inline void add(int u,int v) {
	to[++cnt]=v; nxt[cnt]=head[u]; head[u]=cnt;
}
void insert(int &y,int l,int r,int v,int c) {
	if(!y) y=++tot; siz[y]+=c;
	if(l==r) return ;
	int mid=(l+r)>>1;
	if(v<=mid) insert(ls[y],l,mid,v,c);
	else insert(rs[y],mid+1,r,v,c);
}
void dfs1(int x,int y) {
	int i; dep[x]=dep[y]+1; fa[x]=y; size[x]=1; dfn[x]=++dfn[0];
	for(i=dfn[x];i<=n;i+=i&(-i)) insert(root[i],0,maxn,val[x],1);
	for(i=head[x];i;i=nxt[i]) if(to[i]!=y) {
		dfs1(to[i],x); size[x]+=size[to[i]]; if(size[to[i]]>size[son[x]]) son[x]=to[i];
	}
	for(i=dfn[0]+1;i<=n;i+=i&(-i)) insert(root[i],0,maxn,val[x],-1);
	out[x]=dfn[0]+1;
}
void dfs2(int x,int t) {
	top[x]=t; if(son[x]) dfs2(son[x],t); int i;
	for(i=head[x];i;i=nxt[i]) if(to[i]!=fa[x]&&to[i]!=son[x]) dfs2(to[i],to[i]);
}
int lca(int x,int y) {
	while(top[x]!=top[y]) {
		if(dep[top[x]]>dep[top[y]]) swap(x,y);
		y=fa[top[y]];
	}
	return dep[x]<dep[y]?x:y;
}
int query(int l,int r,int k) {
	if(l==r) return l;
	int mid=(l+r)>>1,i,sizls=0;
	for(i=1;i<=A[0];i++) sizls+=siz[ls[A[i]]];
	for(i=1;i<=B[0];i++) sizls+=siz[ls[B[i]]];
	for(i=1;i<=C[0];i++) sizls-=siz[ls[C[i]]];
	for(i=1;i<=D[0];i++) sizls-=siz[ls[D[i]]];
	if(k<=sizls) {
		for(i=1;i<=A[0];i++) A[i]=ls[A[i]];
		for(i=1;i<=B[0];i++) B[i]=ls[B[i]];
		for(i=1;i<=C[0];i++) C[i]=ls[C[i]];
		for(i=1;i<=D[0];i++) D[i]=ls[D[i]];
		return query(l,mid,k);
	}else {
		for(i=1;i<=A[0];i++) A[i]=rs[A[i]];
		for(i=1;i<=B[0];i++) B[i]=rs[B[i]];
		for(i=1;i<=C[0];i++) C[i]=rs[C[i]];
		for(i=1;i<=D[0];i++) D[i]=rs[D[i]];
		return query(mid+1,r,k-sizls);
	}
}
int main() {
	scanf("%d%d",&n,&m);
	int i,x,y,k;
	for(i=1;i<=n;i++) scanf("%d",&val[i]);
	for(i=1;i<n;i++) {
		scanf("%d%d",&x,&y); add(x,y); add(y,x);
	}
	dfs1(1,0); dfs2(1,1);
	while(m--) {
		scanf("%d%d%d",&k,&x,&y);
		if(!k) {
			for(i=dfn[x];i<=n;i+=i&(-i)) insert(root[i],0,maxn,val[x],-1);
			for(i=out[x];i<=n;i+=i&(-i)) insert(root[i],0,maxn,val[x],1);
			val[x]=y;
			for(i=dfn[x];i<=n;i+=i&(-i)) insert(root[i],0,maxn,val[x],1);
			for(i=out[x];i<=n;i+=i&(-i)) insert(root[i],0,maxn,val[x],-1);
		}else {
			A[0]=B[0]=C[0]=D[0]=0;
			int l=lca(x,y),len=dep[x]+dep[y]-dep[l]-dep[fa[l]];
			if(len<k) puts("invalid request!");
			else {
				k=len-k+1;
				for(i=dfn[x];i;i-=i&(-i)) A[++A[0]]=root[i];
				for(i=dfn[y];i;i-=i&(-i)) B[++B[0]]=root[i];
				for(i=dfn[l];i;i-=i&(-i)) C[++C[0]]=root[i];
				if(fa[l]) for(i=dfn[fa[l]];i;i-=i&(-i)) D[++D[0]]=root[i];
				printf("%d\n",query(0,maxn,k));
			}
		}
	}
}
/*
5 5
5 1 2 3 4
3 1
2 1
4 3
5 3
2 4 5
0 1 2
2 2 3
2 1 4
3 3 5
*/