SPOJ 10628 Count on a tree （lca+主席树）

题意：给定一棵有n个结点的树，每个点有一个权值，共有m个询问，对于每个询问（u，v，k），回答结点u至v之间第k小的点的权值。

思路：主席树+lca，首先指定一个根结点dfs一次并在此过程中建好主席树，对于对于每个询问，我们只需要考虑四棵树，即T[u], T[v], T[lca(u,v)], 再加上T[fa( lca(u,v) )]，fa( lca(u,v) )表示lca(u, v)的父亲结点。

这样一来问题就和线性序列里第k小的数一样了。

#include<cstdio>
#include<cstring>
#include<cmath>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#include<vector>
#include<map>
#include<queue>
#include<stack>
#include<string>
#include<map>
#include<set>
#include<ctime>
#define eps 1e-6
#define LL long long
#define pii (pair<int, int>)
//#pragma comment(linker, "/STACK:1024000000,1024000000")
using namespace std;

//const int maxn = 100000 + 100;
//const int INF = 0x3f3f3f3f;

const int maxn = 100000+10000;
const int M = 2000000;
int n, q, m, tot;
int t[maxn], w[maxn], fa[maxn];
int T[maxn], lson[M], rson[M], c[M];

struct Quest {
	int l, r, k;
} quest[maxn];

void Init_hash(int k) {
	sort(t, t+k);
	m = unique(t, t+k) - t;
}

int Hash(int x) {
	return lower_bound(t, t+m, x) - t;
}

int build(int l, int r) {
	int root = tot++;
	c[root] = 0;
	if(l != r) {
		int mid = (l+r) >> 1;
		lson[root] = build(l, mid);
		rson[root] = build(mid+1, r);
	}
	return root;
}

int Insert(int root, int pos, int val) {
	int newroot = tot++, tmp = newroot;
	int l = 0, r = m-1;
	c[newroot] = c[root] + val;
    while(l < r) {
        int mid = (l+r)>>1;
        if(pos <= mid) {
            lson[newroot] = tot++; rson[newroot] = rson[root];
            newroot = lson[newroot]; root = lson[root];
            r = mid;
        }
        else {
            rson[newroot] = tot++; lson[newroot] = lson[root];
            newroot = rson[newroot]; root = rson[root];
            l = mid+1;
        }
        c[newroot] = c[root] + val;
    }
	return tmp;
}

int Query(int l_root, int r_root, int lca, int k) {
	int l = 0, r = m -1, lca_root = T[lca], fa_root = fa[lca];
	while(l < r) {
		int mid = (l+r) >> 1;
		int tmp = c[lson[l_root]]+c[lson[r_root]]-c[lson[lca_root]]-c[lson[fa_root]];
		if(tmp >= k) {
			r = mid;
			l_root = lson[l_root]; r_root = lson[r_root]; lca_root = lson[lca_root]; fa_root = lson[fa_root];
		}
		else {
			l = mid + 1;
			k -= tmp;
			l_root = rson[l_root]; r_root = rson[r_root]; lca_root = rson[lca_root]; fa_root = rson[fa_root];
		}
	}
	return l;
}

int pnt[maxn], lca[maxn];
bool vis[maxn];
vector<int> G[maxn], query[maxn], num[maxn];
int find(int x) {
	if(x == pnt[x]) return x;
	return pnt[x] = find(pnt[x]);
}
void dfs_lca(int u) {
	vis[u] = 1; pnt[u] = u;
	int sz1 = G[u].size();
	for(int i = 0; i < sz1; i++) {
		int v = G[u][i];
		if(vis[v]) continue;
		fa[v] = T[u];
		dfs_lca(v);
		pnt[v] = u;
	}
	int sz2 = query[u].size();
	for(int i = 0; i < sz2; i++) {
		int v = query[u][i];
		if(vis[v]) lca[num[u][i]] = find(v);
	}
} 
void init() {
	memset(vis, 0, sizeof(vis));
	for(int i = 1; i <= n; i++) {
		G[i].clear(); 
		query[i].clear(); 
		num[i].clear();
	}
}

void dfs_ZXTree(int cur, int fa) {
	int sz = G[cur].size();
	for(int i = 0; i < sz; i++) {
		int u = G[cur][i];
		if(u == fa) continue;
		T[u] = Insert(T[cur], Hash(w[u]), 1);
		dfs_ZXTree(u, cur);
	}
}

int main() {
    //freopen("input.txt", "r", stdin);
    while(cin >> n >> q) {
		init();
		m = 0; tot = 0;
		for(int i = 1; i <= n; i++) scanf("%d", &w[i]), t[m++] = w[i];
		Init_hash(m);
		build(0, m-1);
		for(int i = 1; i < n; i++) {
			int u, v; scanf("%d%d", &u, &v);
			G[u].push_back(v);
			G[v].push_back(u);
		}
		T[1] = Insert(T[0], Hash(w[1]), 1);
		dfs_ZXTree(1, -1);
		for(int i = 0; i < q; i++) {
			int l, r, k; scanf("%d%d%d", &l, &r, &k);
			quest[i].l = l; quest[i].r = r; quest[i].k = k;
			query[l].push_back(r); query[r].push_back(l);
			num[l].push_back(i); num[r].push_back(i);
		}
		fa[1] = T[0];
		dfs_lca(1);
		for(int i = 0; i < q; i++) {
			printf("%d\n", t[Query(T[quest[i].l], T[quest[i].r], lca[i], quest[i].k)]);
		}
    }
    return 0;
}