【主席树】 SPOJ Count on a tree

给出一颗点权树,然后问a到b路径上的第k大。。。。在树上建立主席树,然后lca搞出线段树,在线段树上查询。。

#include <iostream>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <bitset>
#include <cstdio>
#include <algorithm>
#include <cstring>
#include <climits>
#include <cstdlib>
#include <cmath>
#include <time.h>
#define maxn 100005
#define maxm 200005
#define eps 1e-7
#define mod 1000000007
#define INF 0x3f3f3f3f
#define PI (acos(-1.0))
#define lowbit(x) (x&(-x))
#define mp make_pair
//#define ls o<<1
//#define rs o<<1 | 1
//#define lson o<<1, L, mid 
//#define rson o<<1 | 1, mid+1, R
#define pii pair<int, int>
#pragma comment(linker, "/STACK:16777216")
typedef long long LL;
typedef unsigned long long ULL;
//typedef int LL; 
using namespace std;
//LL qpow(LL a, LL b){LL res=1,base=a;while(b){if(b%2)res=res*base;base=base*base;b/=2;}return res;}
LL powmod(LL a, LL b){LL res=1,base=a;while(b){if(b%2)res=res*base%mod;base=base*base%mod;b/=2;}return res;}
// head

#define lson L, mid
#define rson mid+1, R

struct Edge
{
	int v;
	Edge *next;
}*H[maxn], *edges, E[maxm];

struct node
{
	int cnt;
	node *ch[2];
}*tail, pool[maxm << 4], *h[maxn];

const int M = 20;
int anc[maxn][M];
int dep[maxn];
int f[maxn];
int a[maxn];
int b[maxn];
int c[maxn];
int n, m;

void addedges(int u, int v)
{
	edges->v = v;
	edges->next = H[u];
	H[u] = edges++;
}

node* newnode()
{
	tail->cnt = 0;
	tail->ch[0] = tail->ch[1] = NULL;
	return tail++;
}

void init()
{
	tail = pool;
	edges = E;
	memset(H, 0, sizeof H);
}

int to(int u, int d)
{
	for(int i = M - 1; i >= 0; i--) if(dep[anc[u][i]] >= d) u = anc[u][i];
	return u;
}

int lca(int u, int v)
{
	if(dep[u] < dep[v]) swap(u, v);
	u = to(u, dep[v]);
	for(int i = M - 1; i >= 0; i--) if(anc[u][i] != anc[v][i]) u = anc[u][i], v = anc[v][i];
	return u == v ? u : anc[u][0];
}

void pushup(node *o)
{
	o->cnt = o->ch[0]->cnt + o->ch[1]->cnt;
}

node* update(node *o, int L, int R, int v)
{
	node *p = newnode();
	*p = *o;
	if(L == R) {
		p->cnt++;
		return p;
	}
	int mid = (L + R) >> 1;
	if(v <= mid) p->ch[0] = update(o->ch[0], lson, v);
	else p->ch[1] = update(o->ch[1], rson, v);
	pushup(p);
	return p;
}

void dfs(int u, int fa)
{
	if(u == 1) anc[u][0] = 1;
	else anc[u][0] = fa;
	f[u] = fa;
	h[u] = update(h[fa], 1, n, b[u]);
	for(Edge *e = H[u]; e; e = e->next) if(e->v != fa) {
		int v = e->v;
		dep[v] = dep[u] + 1;
		dfs(v, u);
	}
}

node* build(int L, int R)
{
	if(L == R) return newnode();
	int mid = (L + R) >> 1;
	node *o = newnode();
	o->ch[0] = build(lson);
	o->ch[1] = build(rson);
	return o;
}

int query(node *o1, node *o2, node *o3, node *o4, int L, int R, int k)
{
	if(L == R) return L;
	int mid = (L + R) >> 1;
	int cnt = o1->ch[0]->cnt + o2->ch[0]->cnt - o3->ch[0]->cnt - o4->ch[0]->cnt;
	if(cnt >= k) return query(o1->ch[0], o2->ch[0], o3->ch[0], o4->ch[0], lson, k);
	else return query(o1->ch[1], o2->ch[1], o3->ch[1], o4->ch[1], rson, k - cnt);
}


void work()
{
	int u, v, kk;
	for(int i = 1; i <= n; i++) scanf("%d", &a[i]), c[i] = a[i];
	for(int i = 1; i < n; i++) {
		scanf("%d%d", &u, &v);
		addedges(u, v);
		addedges(v, u);
	}
	sort(c+1, c+n+1);
	int tcnt = unique(c+1, c+n+1) - c - 1;
	for(int i = 1; i <= n; i++) b[i] = lower_bound(c+1, c+tcnt+1, a[i]) - c;

	h[0] = build(1, n);
	dfs(1, 0);
	for(int i = 1; i < M; i++)
		for(int j = 1; j <= n; j++)
			anc[j][i] = anc[anc[j][i-1]][i-1];

	while(m--) {
		scanf("%d%d%d", &u, &v, &kk);
		int t = query(h[u], h[v], h[lca(u, v)], h[f[lca(u, v)]], 1, n, kk);
		printf("%d\n", c[t]);
	}
}

int main()
{
	while(scanf("%d%d", &n, &m)!=EOF) {
		init();
		work();
	}

	return 0;
}


你可能感兴趣的:(主席树)