萌新刚学数据结构1s,求助主席树
查看原帖
萌新刚学数据结构1s,求助主席树
232838
huangkx楼主2022/12/24 17:04

RT,WA on #8,调了一个小时了

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;
const int N = 1e5;
int n, m, q, tot;
int dep[N + 5];
int fa[N + 5][32];
int rt[N + 5];
vector < vector < int > > val(N + 5);
vector < vector < int > > g(N + 5);
vector < int > tmp;
struct Segment_Tree{
	int idx;
	int sum[32 * N + 5], cld[32 * N + 5][2];
	void Initialize()
	{
		idx = 0;
	}
	int Create_Node()
	{
		idx ++;
		sum[idx] = 0, cld[idx][0] = cld[idx][1] = 0;
		return idx;
	}
	void Change(int pu, int & u, int l, int r, int V)
	{
		u = Create_Node();
		sum[u] = sum[pu] + 1;
		if(l == r) return;
		int mid = (l + r) >> 1;
		if(V <= mid) cld[u][1] = cld[pu][1], Change(cld[pu][0], cld[u][0], l, mid, V);
		else cld[u][0] = cld[pu][0], Change(cld[pu][1], cld[u][1], mid + 1, r, V);
	}
	int Query(int u, int v, int l, int r, int K)
	{
		if(l == r) return l;
		int mid = (l + r) >> 1;
		if(sum[cld[v][0]] - sum[cld[u][0]] >= K) return Query(cld[u][0], cld[v][0], l, mid, K);
		else return Query(cld[u][1], cld[v][1], mid + 1, r, K - (sum[cld[v][0]] - sum[cld[u][0]]));
	}
}tr;
void DFS(int u, int pa = 0)
{
	dep[u] = dep[pa] + 1;
	fa[u][0] = pa;
	for(int i = 1; i <= 30; i ++) fa[u][i] = fa[fa[u][i - 1]][i - 1];
	rt[u] = rt[pa];
	for(int i : val[u]){
		tr.Change(rt[u], ++ tot, 1, m, i);
		rt[u] = tot;
	}
	for(int v : g[u]){
		if(v == pa) continue;
		DFS(v, u);
	}
}
int LCA(int u, int v)
{
	if(dep[u] < dep[v]) swap(u, v);
	for(int i = 30; i >= 0; i --) if(dep[fa[u][i]] >= dep[v]) u = fa[u][i];
	if(u == v) return u;
	for(int i = 30; i >= 0; i --) if(dep[fa[u][i]] != dep[fa[v][i]]) u = fa[u][i], v = fa[v][i];
	return fa[u][0];
}
void Query(int u, int v, int a)
{
	int lca = LCA(u, v);
	tmp.clear();
	for(int i = 1; i <= min(a, tr.sum[rt[u]] - tr.sum[rt[lca]]); i ++) tmp.push_back(tr.Query(rt[lca], rt[u], 1, m, i));
	for(int i = 1; i <= min(a, tr.sum[rt[v]] - tr.sum[rt[lca]]); i ++) tmp.push_back(tr.Query(rt[lca], rt[v], 1, m, i));
	for(int i = 1; i <= min(a, (int)val[lca].size()); i ++) tmp.push_back(val[lca][i - 1]);
	sort(tmp.begin(), tmp.end());
	printf("%d ", min(a, (int)tmp.size()));
	for(int i = 1; i <= min(a, (int)tmp.size()); i ++) printf("%d ", tmp[i - 1]);
	puts("");
}
int main()
{
	scanf("%d%d%d", & n, & m, & q);
	for(int i = 1; i <= n - 1; i ++){
		int u, v; scanf("%d%d", & u, & v);
		g[u].push_back(v), g[v].push_back(u);
	}
	for(int i = 1; i <= m; i ++){
		int u; scanf("%d", & u);
		val[u].push_back(i);
	}
	tr.Initialize();
	DFS(1);
	for(int i = 1; i <= q; i ++){
		int u, v, a; scanf("%d%d%d", & u, & v, & a);
		Query(u, v, a);
	}
	return 0;
}
2022/12/24 17:04
加载中...