SPOJ[10628] Count on a tree(LCA+主席树)

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

题意：求树上指定路径的权值第k大。

题解：BFS的时候对每个节点的父节点建立一棵权值线段树，因为当前版本的线段树要依赖于前一个版本，就是其父节点的线段树的信息。

查询的时候就是判断$root[u]+root[v]-root[lca[u,v]]-root[fa[lca(u,v)][0]]$的权值与k的关系。 （root记录节点对应的线段树的编号）

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<stdio.h>
#include<string.h>
#include<queue>
#include<cmath>
#include<map>
#include<set>
#include<vector> 
using namespace std;
#define inf 0x3f3f3f3f
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
#define mem(a,b) memset(a,b,sizeof(a));
#define lowbit(x)  x&-x;  
#define debugint(name,x) printf("%s: %d\n",name,x);
#define debugstring(name,x) printf("%s: %s\n",name,x);
typedef long long ll;
typedef unsigned long long ull;
const double eps = 1e-6;

const int mod = 1e9+7;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
const int maxn = 100010;
const int DEG = 20;
int n,m;
vector<int>ve;
struct node{
    int l,r,sum;
}hjt[maxn*40];
int a[maxn],rt[maxn],cnt;

int getid(int x){return lower_bound(ve.begin(),ve.end(),x)-ve.begin()+1;}
void insert(int l,int r,int pre,int &now,int p){
    hjt[++cnt] = hjt[pre];
    now = cnt;
    hjt[now].sum++;
    if(l == r) return;
    int mid = (l+r)>>1;
    if(p <= mid) insert(l,mid,hjt[pre].l,hjt[now].l,p);
    else insert(mid+1,r,hjt[pre].r,hjt[now].r,p);
}
int query(int l,int r,int u,int v,int lca,int fa,int k){
    if(l == r) return l;
    int tmp = hjt[hjt[u].l].sum+hjt[hjt[v].l].sum-hjt[hjt[lca].l].sum-hjt[hjt[fa].l].sum;//核心
    int mid = (l+r)>>1;
    if(k <= tmp) return query(l,mid,hjt[u].l,hjt[v].l,hjt[lca].l,hjt[fa].l,k);
    else return query(mid+1,r,hjt[u].r,hjt[v].r,hjt[lca].r,hjt[fa].r,k-tmp);
}
struct Edge
{
    int to,next;
}edge[maxn*2];
int head[maxn],tot;
void addedge(int u,int v)
{
    edge[tot].to = v;
    edge[tot].next = head[u];
    head[u] = tot++;
}
int fa[maxn][DEG];
int deg[maxn];

void BFS(int root)
{
    queue<int>que;
    deg[root] = 0;
    fa[root][0] = 0; //根节点的父节点设置为0
    que.push(root);
    insert(1,n,rt[0],rt[root],getid(a[root]));
    while(!que.empty())
    {
        int tmp = que.front();
        que.pop();
        for(int i = 1;i < DEG;i++)
            fa[tmp][i] = fa[fa[tmp][i-1]][i-1];
        for(int i = head[tmp]; i != -1;i = edge[i].next)
        {
            int v = edge[i].to;
            if(v == fa[tmp][0])continue;
            deg[v] = deg[tmp] + 1;
            fa[v][0] = tmp;
            insert(1,n,rt[tmp],rt[v],getid(a[v]));
            que.push(v);
        }

    }
}
int LCA(int u,int v)
{
    if(deg[u] > deg[v])swap(u,v);
    int hu = deg[u], hv = deg[v];
    int tu = u, tv = v;
    for(int det = hv-hu, i = 0; det ;det>>=1, i++)
        if(det&1)
            tv = fa[tv][i];
    if(tu == tv)return tu;
    for(int i = DEG-1; i >= 0; i--)
    {
        if(fa[tu][i] == fa[tv][i])
            continue;
        tu = fa[tu][i];
        tv = fa[tv][i];
    }
    return fa[tu][0];
}
int main() {
    scanf("%d%d",&n,&m);
	for(int i = 1; i <= n; i++)
		ve.push_back(a[i]=read());
	sort(ve.begin(),ve.end());
	ve.erase(unique(ve.begin(),ve.end()),ve.end());
	int u,v;
	mem(head,-1);
	for(int i = 1; i < n; i++){
		u = read(),v = read();
		addedge(u,v);
		addedge(v,u);
	}
	BFS(1);
	int k;
	while(m--){
		scanf("%d%d%d",&u,&v,&k);
		int lca = LCA(u,v);
		printf("%d\n",ve[query(1,n,rt[u],rt[v],rt[lca],rt[fa[lca][0]],k)-1]);
	}
}