Codeforces Round #326 (Div. 1) C - Duff in the Army

Posted 2020-12-16 lordxx

树上主席树。
题意要求求u-v的路径上，前k小个数字是多少。转换一下，每次求1-k小的数字。故我们想到了树上主席树。
基于树上点前缀和的特点有：(sum[u]+sum[v]-sum[lca]-sum[fa[lca]])是u-v路径的和。我们基于此，使用主席树，每个子树是权值线段树。就可以使用前缀和的思想来求得第k小。

#include<iostream>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<cstdlib>
#include<climits>
#include<stack>
#include<vector>
#include<queue>
#include<set>
#include<bitset>
#include<map>
//#include<regex>
#include<cstdio>
#include <iomanip>
#pragma GCC optimize(2)
#define up(i,a,b)  for(int i=a;i<b;i++)
#define dw(i,a,b)  for(int i=a;i>b;i--)
#define upd(i,a,b) for(int i=a;i<=b;i++)
#define dwd(i,a,b) for(int i=a;i>=b;i--)
//#define local
typedef long long ll;
typedef unsigned long long ull;
const double esp = 1e-6;
const double pi = acos(-1.0);
const int INF = 0x3f3f3f3f;
const int inf = 1e9;
using namespace std;
ll read()
{
	char ch = getchar(); ll x = 0, f = 1;
	while (ch<‘0‘ || ch>‘9‘) { if (ch == ‘-‘)f = -1; ch = getchar(); }
	while (ch >= ‘0‘ && ch <= ‘9‘) { x = x * 10 + ch - ‘0‘; ch = getchar(); }
	return x * f;
}
typedef pair<int, int> pir;
#define lson l,mid,root<<1
#define rson mid+1,r,root<<1|1
#define lrt root<<1
#define rrt root<<1|1
int n, m, q;
const int N = 1e5 + 10;
vector<int>vec[N];
vector<int>val[N];
int dp[2 * N][30];
int dep[N * 2];
int root[N];
int ver[2 * N];
int cnt = 0;
int par[N];
int rk[N];
struct zx {
	int tot;
	int tree[N * 20];
	int ls[N * 20];
	int rs[N * 20];
	void update(int &o, int pre, int l, int r, int pos)
	{
		o = ++tot;
		ls[o] = ls[pre], rs[o] = rs[pre];
		tree[o] = tree[pre] + 1;
		if (l == r)return;
		int mid = (l + r) >> 1;
		if (pos <= mid)update(ls[o], ls[pre], l, mid, pos);
		else update(rs[o], rs[pre], mid + 1, r, pos);
	}
	int query(int lca,int flca,int L,int R, int l, int r, int k)
	{
		if (tree[L] + tree[R] - tree[lca] - tree[flca] < k)return -1;
		if (l == r)return l;
		int mid = (l + r) >> 1;
		int temp = tree[ls[L]] + tree[ls[R]] - tree[ls[lca]] - tree[ls[flca]];
		if (k <= temp)return query(ls[lca], ls[flca], ls[L], ls[R], l, mid, k);
		else return query(rs[lca], rs[flca], rs[L], rs[R], mid + 1, r, k - temp);
	}
}T;
void dfs(int u, int fa,int d)
{
	ver[++cnt] = u;
	dep[cnt] = d;
	rk[u] = cnt;
	root[u] = root[fa];
	par[u] = fa;
	for (auto k : val[u])
	{
		T.update(root[u], root[u], 1, m, k);
	}
	for (auto k : vec[u])
	{
		if (k == fa)continue;
		dfs(k, u, d + 1);
		ver[++cnt] = u;
		dep[cnt] = d;
	}
}
void init()
{
	upd(i, 0, cnt)dp[i][0] = i;
	for (int j = 1; j <= 30; j++)
	{
		for (int i = 1; i + (1 << j) - 1 <= 2 * n; i++)
		{
			int t1 = dp[i][j - 1];
			int t2 = dp[i + (1 << (j - 1))][j - 1];
			dp[i][j] = dep[t1] > dep[t2] ? t2 : t1;
		}
	}
}
int stquery(int l, int r)
{
	l = rk[l], r = rk[r];
	if (l > r)swap(l, r);
	int k = int(log2(r - l + 1));
	int t1 = dp[l][k];
	int t2 = dp[r - (1 << k) + 1][k];
	return dep[t1] < dep[t2] ? ver[t1] : ver[t2];
}
int main()
{
	n = read(), m = read(), q = read();
	int u, v;
	upd(i, 1, n - 1)
	{
		u = read(), v = read();
		vec[u].push_back(v);
		vec[v].push_back(u);
	}
	upd(i, 1, m)
	{
		u = read();
		val[u].push_back(i);
	}
	dfs(1, 0, 1);
	init();
	int p;
	while (q--)
	{
		u = read(), v = read(), p = read();
		int lca = stquery(u, v);
		int flca = par[lca];
		vector<int>ans;
		upd(i, 1, p)
		{
			int temp = T.query(root[lca], root[flca], root[u], root[v], 1, m, i);
			if (temp == -1)break;
			ans.push_back(temp);
		}
		printf("%d ", ans.size());
		for (auto k : ans)
		{
			printf("%d ", k);
		}	
		printf("
");
	}
	return 0;
}