SPOJ 10628 Count on a tree【主席树+LCA】

$$Countonatree$$

【https://vjudge.net/problem/SPOJ-COT 】
Description：
You are given a tree with N nodes. The tree nodes are numbered from 1 to N. Each node has an integer weight.

We will ask you to perform the following operation:

u v k : ask for the kth minimum weight on the path from node u to node v
Input
In the first line there are two integers N and M. (N, M <= 100000)

In the second line there are N integers. The ith integer denotes the weight of the ith node.

In the next N-1 lines, each line contains two integers u v, which describes an edge (u, v).

In the next M lines, each line contains three integers u v k, which means an operation asking for the kth minimum weight on the path from node u to node v.

Output
For each operation, print its result.

Example
Input:
8 5
105 2 9 3 8 5 7 7
1 2
1 3
1 4
3 5
3 6
3 7
4 8
2 5 1
2 5 2
2 5 3
2 5 4
7 8 2
Output:
2
8
9
105
7

TJ：

其实就是在树上的主席树，普通的主席树查询区间[L,R]第K大的时候建树是按区间[0,i]的权值建树，而树上的主席树则按根节点到当前点的权值建树，当查询节点u~v之间的第k大时，和普通主席树的求法差不多，还是用类似前缀和的方法来做，因为每棵树都是记录根节点到当前节点的权值，所以减去重复计算的时候要用到LCA，记u和v的LCA为lca
则u到v区间的权值和表示为：
sum = node[u].sum+node[v].sum-node[lca].sum-node[fa[lca]].sum
然后和普通主席树一样做就好了

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
#include<cmath>
using namespace std;
const int maxn = 1e5+7;
vector<int> vec;
int getid(int x){ return lower_bound(vec.begin(),vec.end(),x)-vec.begin()+1; }
int n,m,seq[maxn],root[maxn],fa[maxn][19],dep[maxn],head[maxn],edge_num,maxdep;
struct EDGE{
	int to,next;
}edge[maxn<<2];
struct ljt_Tree{
	struct Node{
		int l,r,sum;
	}node[maxn<<5];
	int cnt;
	ljt_Tree(){cnt=0;}
	void update(int l,int r,int &now,int pre,int pos){
		node[++cnt] = node[pre];
		node[now=cnt].sum++;
		if(l+1==r) return;
		int mid = (l+r)>>1;
		if(pos<mid) update(l,mid,node[now].l,node[pre].l,pos);
		else update(mid,r,node[now].r,node[pre].r,pos);
	}
	int query(int l,int r,int x,int y,int lca,int flca,int k){
		if(l+1==r) return l;
		int sum = node[node[x].l].sum+node[node[y].l].sum-node[node[lca].l].sum-node[node[flca].l].sum;
		int mid = (l+r)>>1;
		if(sum>=k) return query(l,mid,node[x].l,node[y].l,node[lca].l,node[flca].l,k);
		else return query(mid,r,node[x].r,node[y].r,node[lca].r,node[flca].r,k-sum);
	}
}Seg_Tree;
void add_edge(int u,int v){
	edge[edge_num].to = v;
	edge[edge_num].next = head[u];
	head[u] = edge_num++;
}
void dfs(int now,int pre,int depth){
	maxdep = max(maxdep,depth);
	dep[now] = depth;
	if(now) Seg_Tree.update(1,n+1,root[now],root[pre],getid(seq[now]));
	fa[now][0] = pre;
	for(int i=1;i<=18;i++) fa[now][i] = fa[fa[now][i-1]][i-1];
	for(int i=head[now];i!=-1;i=edge[i].next){
		int v = edge[i].to;
		if(v==pre) continue;
		dfs(v,now,depth+1);
	}
}
int LCA(int u,int v){
	if(dep[u]<dep[v]) swap(u,v);
	for(int i=0;dep[u]!=dep[v];i++) if((dep[u]-dep[v])>>i & 1) u = fa[u][i];
	if(u==v) return u;
	for(int i=log2(maxdep);i>=0;i--){
		if(fa[u][i]!=fa[v][i]){
			u = fa[u][i];
			v = fa[v][i];
		}
	}
	return fa[u][0];
}
int main(){
	scanf("%d %d",&n,&m);
	memset(head,255,sizeof(head));
	for(int i=1;i<=n;i++) scanf("%d",&seq[i]),vec.push_back(seq[i]);
	sort(vec.begin(),vec.end());
	vec.erase(unique(vec.begin(),vec.end()),vec.end());
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d %d",&u,&v);
		add_edge(u,v),add_edge(v,u);
	}
	add_edge(0,1),add_edge(1,0);
	dfs(0,0,0);
	for(int i=1;i<=m;i++){
		int u,v,k;
		scanf("%d %d %d",&u,&v,&k);
		int lca = LCA(u,v);
		printf("%d\n",vec[Seg_Tree.query(1,n+1,root[v],root[u],root[lca],root[fa[lca][0]],k)-1]);
	}
	return 0;
}