P2633 Count on a tree

Description

题目描述
给定一棵N个节点的树，每个点有一个权值，对于M个询问(u,v,k)，你需要回答u xor lastans和v这两个节点间第K小的点权。其中lastans是上一个询问的答案，初始为0，即第一个询问的u是明文。

Solution

按照dfs序建主席树, 那么l, r的答案就是(R_l+R_r-R_{lca}-R_{lca_fa}).

Code

#include <stdio.h>
#include <string.h>
#include <iostream>
#include <algorithm>
const int N = 100005;

struct Edge {
	int v, nxt;
} e[N << 1];
int head[N], tot;
void AddEdge(int u, int v) {
	e[++tot] = (Edge) {v, head[u]}; head[u] = tot;
	e[++tot] = (Edge) {u, head[v]}; head[v] = tot;
}

struct Node {
	int val;
	Node *ls, *rs;
	Node() { ls = rs = NULL, val = 0; }
	Node(Node* co) : Node() { val = co->val; }
	void update() { val = ls->val + rs->val; }
	void clear() { if (ls) ls->clear(), rs->clear(); delete this; }
};

class Tree {
private:
	int n;
	std:: vector<Node*> R;
	void build(int l, int r, Node* node) {
		if (l == r) return;
		int mid = l + r >> 1;
		node->ls = new Node();
		node->rs = new Node();
		build(l, mid, node->ls);
		build(mid + 1, r, node->rs);
	}
	void rebuild(int l, int r, int p, Node* node, Node* pre) {
		if (l == r) {
			node->val += 1;
			return;
		}
		int mid = l + r >> 1;
		if (p <= mid) {
			node->rs = pre->rs;
			node->ls = new Node(pre->ls);
			rebuild(l, mid, p, node->ls, pre->ls);
		} else if (p > mid) {
			node->ls = pre->ls;
			node->rs = new Node(pre->rs);
			rebuild(mid + 1, r, p, node->rs, pre->rs);
		}
		node->update();
	}
	Node* merge(Node *l, Node *r) {
		if (not l and not r) return NULL;
		if (not l) return r; if (not r) return l;
		Node* res = new Node();
		res->val = l->val + r->val;
		res->ls = merge(l->ls, r->ls);
		res->rs = merge(l->rs, r->rs);
		return res;
	}
	int kth_query(int l, int r, int k, Node *a, Node *b) {
		if (l == r) return l;
		int mid = l + r >> 1;
		int val = b->ls->val - a->ls->val;
		if (val < k)
			return kth_query(mid + 1, r, k - val, a->rs, b->rs);
		else return kth_query(l, mid, k, a->ls, b->ls);
	}
public:
	void build(int _) {
		for (int i = 0; i < _; i += 1)
			R.push_back(new Node());
		n = _; build(1, n, R[0]);
	}
	void rebuild(int p, int now, int fa) {
		rebuild(1, n, p, R[now], R[fa]);
	}
	int get(int l, int r, int ll, int rr, int p) {
		Node *he1 = merge(R[l], R[r]), *he2 = merge(R[ll], R[rr]);
		return kth_query(1, n, p, he1, he2);
	}
};

int f[N][20], dep[N];
int val[N], A[N];
int cnt;
#define getval(X) std:: lower_bound(A + 1, A + cnt + 1, X) - A
void dfs(int u, int fa, Tree* T) {
	T->rebuild(getval(u), u, fa);
	f[u][0] = fa; dep[u] = dep[fa] + 1;
	for (int i = head[u]; i; i = e[i].nxt) {
		if (u == fa) continue;
		dfs(e[i].v, u, T);
	}
}
int lca(int u, int v) {
	if (dep[u] < dep[v]) std:: swap(u, v);
	for (int i = 19; ~i; i -= 1)
		if (dep[f[u][i]] >= dep[v]) 
			u = f[u][i];
	for (int i = 19; ~i; i -= 1)
		if (f[u][i] != f[v][i]) 
			u = f[u][i], v = f[v][i];
	return u == v ? u : f[u][0];
}

int main () {
	int n, m;
	scanf("%d%d", &n, &m); 
	for (int i = 1; i <= n; i += 1)
		scanf("%d", &val[i]), A[i] = val[i];
	std:: sort(A + 1, A + n + 1);
	cnt = std:: unique(A + 1, A + n + 1) - A - 1;
	for (int i = 1; i < n; i += 1) {
		int u, v; scanf("%d%d", &u, &v);
		AddEdge(u, v);
	}
	Tree* T = new Tree();
	T->build(n);
	dfs(1, 0, T);
	for (int i = 1; i < 20; i += 1)
		for (int j = 1; j <= n; j += 1)
			f[j][i] = f[f[j][i - 1]][i - 1];
	int lastans = 0;
	for (int i = 0; i < m; i += 1) {
		int u, v, w, Lca; scanf("%d%d%d", &u, &v, &w); 
		u ^= lastans; Lca = lca(u, v);
		printf("%d
", A[T->get(u, v, Lca, f[Lca][0], w)]);
	}
	return 0;
}