[树上主席树][洛谷 P2633] Count on a tree

#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdio>
#include <vector>
using namespace std;

#define RG register int
#define LL long long

template<typename elemType>
inline void Read(elemType &T){
    elemType X=0,w=0; char ch=0;
    while(!isdigit(ch)) {w|=ch=='-';ch=getchar();}
    while(isdigit(ch)) X=(X<<3)+(X<<1)+(ch^48),ch=getchar();
    T=(w?-X:X);
}

const int INF=2147483647;
const int MAXN=100010;
vector<int> HashData;
int Data[MAXN];
struct edge{int Next,to;};
edge G[200010];
int head[100010],Value[100010],Deep[100010],Anc[100010][18];
int N,M,cnt=2,ElemNum;

inline void Discretization(){
    HashData[0]=-INF;
    sort(HashData.begin(),HashData.end());
    HashData.erase(unique(HashData.begin(),HashData.end()),HashData.end());
    ElemNum=HashData.size()-1;
    return;
}

inline int GetPos(int x){
    return lower_bound(HashData.begin(),HashData.end(),x)-HashData.begin();
}

struct SegmentTree{
    struct SegTreeNode{
        int Size;
        int Lson,Rson;
    };
    SegTreeNode SegTree[MAXN<<5];
    int ROOT[MAXN];
    int Index=0;

    inline void Push_Up(int Root){
        int ls=SegTree[Root].Lson;
        int rs=SegTree[Root].Rson;
        SegTree[Root].Size=SegTree[ls].Size+SegTree[rs].Size;
        return;
    }

    void Build_Empty_SegTree(int &Root,int L,int R){
        Root=++Index;
        if(L==R){SegTree[Root].Size=0;return;}
        int mid=(L+R)>>1;
        Build_Empty_SegTree(SegTree[Root].Lson,L,mid);
        Build_Empty_SegTree(SegTree[Root].Rson,mid+1,R);
        Push_Up(Root);
        return;
    }

    void Insert_Node(int &Pre,int &Cur,int L,int R,int Pos){
        Cur=++Index;
        SegTree[Cur]=SegTree[Pre];
        if(L==R){++SegTree[Cur].Size;return;}
        int mid=(L+R)>>1;
        if(Pos<=mid) Insert_Node(SegTree[Pre].Lson,SegTree[Cur].Lson,L,mid,Pos);
        else Insert_Node(SegTree[Pre].Rson,SegTree[Cur].Rson,mid+1,R,Pos);
        Push_Up(Cur);
        return;
    }

    int Get_Kth(int fa,int lca,int u,int v,int L,int R,int Kth){
        if(L==R) return L;
        int mid=(L+R)>>1;
        int UL=SegTree[u].Lson,VL=SegTree[v].Lson;
        int FaL=SegTree[fa].Lson,LcaL=SegTree[lca].Lson;
        int Lsize=SegTree[UL].Size+SegTree[VL].Size-SegTree[FaL].Size-SegTree[LcaL].Size;
        if(Kth<=Lsize)
            return Get_Kth(SegTree[fa].Lson,SegTree[lca].Lson,SegTree[u].Lson,SegTree[v].Lson,L,mid,Kth);
        return Get_Kth(SegTree[fa].Rson,SegTree[lca].Rson,SegTree[u].Rson,SegTree[v].Rson,mid+1,R,Kth-Lsize);
    }
};
SegmentTree Tree;

inline void Add_Edge(int u,int v){
    G[cnt].to=v;
    G[cnt].Next=head[u];
    head[u]=cnt++;
}

void DFS_Init(int u,int fa){
    Tree.Insert_Node(Tree.ROOT[fa],Tree.ROOT[u],1,ElemNum,GetPos(Value[u]));
	Anc[u][0]=fa;
	for(int i=1;i<18;++i)
		Anc[u][i]=Anc[Anc[u][i-1]][i-1];
	for(int i=head[u];i;i=G[i].Next){
		int v=G[i].to;
		if(v==fa) continue;
		Deep[v]=Deep[u]+1;
		DFS_Init(v,u);
	}
	return;
}

int LCA(int u,int v){
	if(u==1 || v==1) return 1;
	if(Deep[u]<Deep[v]) swap(u,v);
	for(RG i=17;i>=0;--i){
		if(Deep[Anc[u][i]]>=Deep[v])
			u=Anc[u][i];
	}
	if(u==v) return u;
	for(RG i=17;i>=0;--i){
		if(Anc[u][i]!=Anc[v][i]){
			u=Anc[u][i];	
			v=Anc[v][i];
		}
	} 
	return Anc[u][0];
}

int main(){
    Read(N);Read(M);
    HashData.resize(N+1);
    for(RG i=1;i<=N;++i){
        Read(Value[i]);
        HashData[i]=Value[i];
    }
    for(RG i=1;i<=N-1;++i){
        int u,v;
        Read(u);Read(v);
        Add_Edge(u,v);
        Add_Edge(v,u);
    }
    Discretization();
    Tree.Build_Empty_SegTree(Tree.ROOT[0],1,ElemNum);
    DFS_Init(1,0);
    int LastAns=0;
    while(M--){
        int u,v,Kth,Pos,lca;
        Read(u);Read(v);Read(Kth);
        u=u^LastAns;lca=LCA(u,v);
        Pos=Tree.Get_Kth(Tree.ROOT[Anc[lca][0]],Tree.ROOT[lca],Tree.ROOT[u],Tree.ROOT[v],1,ElemNum,Kth);
        printf("%d\n",HashData[Pos]);
        LastAns=HashData[Pos];
    }
    return 0;
}

猜你喜欢

转载自www.cnblogs.com/AEMShana/p/12664104.html