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