P5838 [USACO19DEC]Milk Visits G

题解

并不需要什么高级数据结构
用树链剖分维护
对于每种颜色开个 ( ext{vector})，然后把是这种颜色的点的 ( ext{dfs}) 序加进来排序
对于 ([dfn[top[x]],dfn[x]]) 这一区间问有没有某种颜色
相当于问某种颜色有没有至少一个在这个区间内
直接二分查 ( ext{vector}) 即可

(Code)

#include<cstdio>
#include<vector>
#include<algorithm>
using namespace std;

const int N = 1e5 + 5;
int n, m, a[N], h[N], ans[N];
vector<int> col[N];

struct edge{int to, nxt;}e[N << 1];
inline void add(int x, int y)
{
	static int tot = 0;
	e[++tot] = edge{y, h[x]}, h[x] = tot;
}

int siz[N], son[N], dfn[N], fa[N], top[N], dep[N];
void dfs1(int x)
{
	siz[x] = 1, dep[x] = dep[fa[x]] + 1;
	for(register int i = h[x]; i; i = e[i].nxt)
	{
		int v = e[i].to;
		if (v == fa[x]) continue;
		fa[v] = x, dfs1(v), siz[x] += siz[v];
		if (siz[v] > siz[son[x]]) son[x] = v;
	}
}
void dfs2(int x)
{
	static int dfc = 0;
	dfn[x] = ++dfc;
	if (son[x]) top[son[x]] = top[x], dfs2(son[x]);
	for(register int i = h[x]; i; i = e[i].nxt)
	{
		int v = e[i].to;
		if (v == fa[x] || v == son[x]) continue;
		top[v] = v, dfs2(v);
	}
}

inline int search(int x, int y, int c)
{
	int l = 0, r = col[c].size() - 1, mid, ll = -1, rr = -1;
	while (l <= r)
	{
		mid = (l + r) >> 1;
		if (col[c][mid] >= x) ll = mid, r = mid - 1;
		else l = mid + 1;
	}
	if (ll == -1) return 0;
	l = ll, r = col[c].size() - 1;
	while (l <= r)
	{
		mid = (l + r) >> 1;
		if (col[c][mid] <= y) rr = mid, l = mid + 1;
		else r = mid - 1;
	}
	if (rr == -1) return 0;
	return 1;
}

inline int solve(int x, int y, int c)
{
	int fx = top[x], fy = top[y];
	while (fx ^ fy)
	{
		if (dep[fx] > dep[fy])	
		{
			if (search(dfn[fx], dfn[x], c)) return 1;
			x = fa[fx], fx = top[x];
		}
		else{
			if (search(dfn[fy], dfn[y], c)) return 1;
			y = fa[fy], fy = top[y];
		}
	}
	if (dep[x] < dep[y]) return search(dfn[x], dfn[y], c);
	return search(dfn[y], dfn[x], c);
}

int main()
{
	scanf("%d%d", &n, &m);
	for(register int i = 1; i <= n; i++) scanf("%d", a + i);
	for(register int i = 1, u, v; i < n; i++) scanf("%d%d", &u, &v), add(u, v), add(v, u);
	dfs1(1), top[1] = 1, dfs2(1);
	for(register int i = 1; i <= n; i++) col[a[i]].push_back(dfn[i]);
	for(register int i = 1; i <= n; i++) sort(col[i].begin(), col[i].end());
	int cnt = 0;
	for(int u, v, w; m; --m)
	{
		scanf("%d%d%d", &u, &v, &w);
		if (solve(u, v, w)) ans[++cnt] = 1;
		else ans[++cnt] = 0;
	}
	for(register int i = 1; i <= cnt; i++) printf("%d", ans[i]);
}