SP1487 PT07J - Query on a tree III

题意

给定一颗树，树有点权。多次询问，每次询问以 $u$ 为根的子树中第 $k$ 小的点权的点编号，点权互不相同。

解法

要求树上询问问题，可以想到树剖套各种数据结构，但是注意到这道题只需要求子树的问题，所以考虑用 dfs 序和主席树解决这个问题。

主席树其实就是【模板】可持久化线段树 2，dfs 序很好处理，所以这道题很容易地被切掉了！

另外注意，询问求的是点的编号，而非点权。

代码：

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <vector>
#include <unordered_map>
using namespace std;

const int N = 1e5 + 5;

int a[N], n, m, na[N];
unordered_map<int, int> mp;
vector<int> G[N], b;

class Persis_Segment_Tree
{
public:
	struct Node
	{
		int l, r, sum;
	}tr[N * 23];
	int root[N], idx;
	Persis_Segment_Tree() : idx(0) {}
	int build(int l, int r)
	{
		int q = ++idx;
		tr[q].sum = 0;
		if (l == r) return q;
		int mid = l + r >> 1;
		tr[q].l = build(l, mid);
		tr[q].r = build(mid + 1, r);
		return q;
	}
	int update(int u, int l, int r, int x)
	{
		int q = ++idx;
		tr[q] = tr[u];
		if (l == x and r == x)
		{
			tr[q].sum++;
			return q;
		}
		int mid = l + r >> 1;
		if (x <= mid) tr[q].l = update(tr[u].l, l, mid, x);
		else tr[q].r = update(tr[u].r, mid + 1, r, x);
		tr[q].sum = tr[tr[q].l].sum + tr[tr[q].r].sum;
		return q;
	}
	int query(int old, int now, int l, int r, int k)
	{
		if (l == r) return mp[b[l]];
		int cnt = tr[tr[now].l].sum - tr[tr[old].l].sum, mid = l + r >> 1;
		if (k <= cnt)
		{
			return query(tr[old].l, tr[now].l, l, mid, k);
		}
		return query(tr[old].r, tr[now].r, mid + 1, r, k - cnt);
	}
};

class TreeCut
{
public:
	int idx, id[N], sz[N];
	Persis_Segment_Tree sg;
	TreeCut() : idx(0) {}
	void dfs(int u, int father)
	{
		id[u] = ++idx;
		sz[u] = 1;
		na[idx] = a[u];
		for (int i = 0; i < G[u].size(); i++)
		{
			int nx = G[u][i];
			if (nx != father)
			{
				dfs(nx, u);
				sz[u] += sz[nx];
			}
		}
	}
	void build()
	{
		dfs(1, 1);
		for (int i = 1; i <= n; i++) b.push_back(na[i]);
		sort(b.begin(), b.end());
		//b.erase(unique(b.begin(), b.end()), b.end());
		sg.root[0] = sg.build(0, b.size() - 1);
		for (int i = 1; i <= n; i++)
		{
			int g = lower_bound(b.begin(), b.end(), na[i]) - b.begin();
			sg.root[i] = sg.update(sg.root[i - 1], 0, b.size() - 1, g);
		}
	}
	int query(int u, int k)
	{
		int l = id[u], r = id[u] + sz[u] - 1;
		return sg.query(sg.root[l - 1], sg.root[r], 0, b.size() - 1, k);
	}
};

TreeCut tc;

int main()
{
	scanf("%d", &n);
	for (int i = 1; i <= n; i++)
	{
		scanf("%d", &a[i]);
		mp[a[i]] = i;
	}
	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);
	}
	tc.build();
	scanf("%d", &m);
	while (m--)
	{
		int u, k;
		scanf("%d%d", &u, &k);
		printf("%d\n", tc.query(u, k));
	}
	return 0;
}