SPOJ COT 10628 Count on a tree

这题和普通的第K大类似。

普通的第K大，是从后往前建立主席树，前一个在后一个的基础上修改。

而树上第K大，依旧是每个结点一棵主席树，是在父节点的基础上修改。

这里用了倍增法求lca

这代码要跑2900MS左右。。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <queue>
#include <vector>
#include <algorithm>
using namespace std;
const int N = 200010;
const int M = N * 40;
int n,m,q,tot;
int a[N],t[N];
int T[N],ls[M],rs[M],c[M],d[N],fa[N][27];
vector<int> vec[N];
void init(){
    memset(fa, -1, sizeof(fa));
    tot = 0;
}
void Init_hash(){
	for(int i = 1; i<= n; i++){
		t[i] = a[i];
	}
	sort(t+1,t+n+1);
	m = unique(t+1,t+1+n)-t-1;
}
void build(int l, int r, int &rt){
	rt = tot ++;
	c[rt] = 0;
	if(l != r){
		int mid = (l+r)>>1;
		build(l, mid, ls[rt] );
		build(mid+1, r, rs[rt]);
	}
}
int hash(int x){
	return lower_bound(t+1,t+1+m,x)-t;
}
int update(int root, int pos, int val){//循环
	int newroot = tot++, tmp = newroot;
	c[newroot] = c[root] + val;
	int l = 1, r = m;
	while(l < r){
		int mid = (l+r)>>1;
		if(pos <= mid){
			ls[newroot] = tot++;
			rs[newroot] = rs[root];
			newroot = ls[newroot];
			root = ls[root];
			r = mid;
		}
		else {
			rs[newroot] = tot++;
			ls[newroot] = ls[root];
			newroot = rs[newroot];
			root = rs[root];
			l = mid+1;
		}
		c[newroot] = c[root] + val;
	}
	return tmp;
}
void update2(int last, int pos, int l, int r, int& rt){//递归
    rt = ++ tot;
    ls[rt] = ls[last], rs[rt] = rs[last], c[rt] = c[last] + 1;
    if (l == r) return ;
    int mid = (l + r) >> 1;
    if (pos <= mid) update2(ls[last], pos, l, mid, ls[rt]);
    else update2(rs[last], pos, mid + 1, r, rs[rt]);
}
int query(int lrt, int rrt, int lca, int k){//循环
	int lcart = T[lca], pos = hash(a[lca]);
	int l = 1, r = m;
	while(l < r){
		int mid = (l+r)>>1;
		int cot = c[ls[lrt]] + c[ls[rrt]] - 2*c[ls[lcart]] + (pos >= l && pos <=mid);
		if(cot >= k){
			r = mid;
			lrt = ls[lrt];
			rrt = ls[rrt];
			lcart = ls[lcart];
		}
		else {
			l = mid+1;
			k -= cot;
			lrt = rs[lrt];
			rrt = rs[rrt];
			lcart = rs[lcart];
		}
	}
	return l;
}
int query2(int pos, int left_rt, int right_rt, int lca_rt, int l, int r, int k){//递归
    if (l == r) return l;
    int mid = (l + r) >> 1;
    int cnt = c[ls[left_rt]] + c[ls[right_rt]] - 2 * c[ls[lca_rt]] + (pos >= l && pos <= mid);
    if (k <= cnt) return query2(pos, ls[left_rt], ls[right_rt], ls[lca_rt], l, mid, k);
    else return query2(pos, rs[left_rt], rs[right_rt], rs[lca_rt], mid + 1, r, k - cnt);
}
void dfs_build(int u, int pre){
	fa[u][0] = pre;
	d[u] = d[pre] + 1;
    int pos = hash(a[u]);
    T[u] = update(T[pre], pos, 1);
    //update2(T[pre], pos, 1, m, T[u]);
    int sz = vec[u].size();
    for(int i = 0; i < sz; i++){
        int v = vec[u][i];
        if(v != pre){
            dfs_build(v, u);
        }
    }
}
void init_fa(){
    for(int j = 1; (1<<j) <= n; j++){
        for(int i = 1; i <= n; i++){
            fa[i][j] = fa[fa[i][j-1]][j-1];
        }
    }
}
int LCA(int a, int b){
    int i, j;
    if(d[a] < d[b])swap(a, b);
    for(i = 0; (1<<i) <= d[a]; i++);i--;
    for( j = i; j >= 0; j--){
        if(d[a] - (1<<j) >= d[b])
            a = fa[a][j];
    }
    if(a == b)return a;
    for(j = i; j >= 0; j--){
        if(fa[a][j] != -1 && fa[a][j] != fa[b][j]){
            a = fa[a][j], b = fa[b][j];
        }
    }
    return fa[a][0];
}

int main(){
	int u,v,k;
	//freopen("in.txt","r",stdin);
	while(~scanf("%d%d", &n, &q)){
		init();
		for(int i = 1; i <= n; i++ )scanf("%d", a+i);
		for(int i = 1; i < n; i++){
			scanf("%d%d", &u, &v);
			vec[u].push_back(v);
			vec[v].push_back(u);
		}
		Init_hash();
		build(1, m, T[0]);
		dfs_build(1, 0);
		init_fa();
		while(q--){
			scanf("%d%d%d", &u, &v, &k);
			int lca = LCA(u, v);
			printf("%d\n",t[query(T[u],T[v],lca,k)]);
			//printf("%d\n",t[query2(hash(a[lca]), T[u], T[v], T[lca],1, m, k)]);
		}
		for(int i = 1; i <= n; i++)vec[i].clear();
	}
	return 0;
}