【[USACO19DEC】Milk Visits G - 重链剖分

题目描述

Farmer John 计划建造 N 个农场，用 N-1 条道路连接，构成一棵树。每个农场有一头奶牛，品种为 1 到 N 之间的一个整数 Ti

Farmer John 的 M 个朋友经常前来拜访他。在朋友 i 拜访之时，Farmer John 会与他的朋友沿着从农场 Ai 到农场 Bi 之间的唯一路径行走（可能有 Ai = Bi）。除此之外，他们还可以品尝他们经过的路径上任意一头奶牛的牛奶。由于 Farmer John 的朋友们大多数也是农场主，他们对牛奶有着极强的偏好。他的每个朋友都只喝某种特定品种的奶牛的牛奶。任何 Farmer John 的朋友只有在他们访问时能喝到他们偏好的牛奶才会高兴。

请求出每个朋友在拜访过后是否会高兴。

题解

将树重链剖分，每个颜色开一个 vector，将同一颜色的点的 dfn 加入 vector

每次向上跳 top 时，lower_bound 一下 dfn[top[u]]，若小于 dfn[u] 则在这条重链上有答案

//I am hunting for the golden stag.
#include <algorithm>
#include <vector>
#include <cstdio>
using namespace std;
const int maxn = 1e5 + 10;
int n,m,cnt,nod[maxn],dep[maxn],top[maxn],dfn[maxn],fa[maxn],siz[maxn],son[maxn];
vector<int> edge[maxn],col[maxn];
inline void dfs1(int now,int f) {
	dep[now] = dep[f]+1;
	fa[now] = f;
	siz[now] = 1;
	for (size_t i = 0;i < edge[now].size();i++) {
		int to = edge[now][i];
		if (to ^ f) {
			dfs1(to,now);
			siz[now] += siz[to];
			if (siz[to] > siz[son[now]]) son[now] = to;
		}
	}
}
inline void dfs2(int now,int ntop) {
	dfn[now] = ++cnt;
	top[now] = ntop;
	col[nod[now]].push_back(cnt);
	if (son[now]) dfs2(son[now],ntop);
	for (size_t i = 0;i < edge[now].size();i++) {
		int to = edge[now][i];
		if (to ^ fa[now] && to ^ son[now]) dfs2(to,to);
	}
}
inline int query(int u,int v,int w) {
	for (;top[u] ^ top[v];u = fa[top[u]]) {
		if (dep[top[u]] < dep[top[v]]) swap(u,v);
		size_t num = lower_bound(col[w].begin(),col[w].end(),dfn[top[u]])-col[w].begin();
 		if (num < col[w].size() && col[w][num] <= dfn[u]) return 1;
	}
	if (dep[u] > dep[v]) swap(u,v);
	size_t num = lower_bound(col[w].begin(),col[w].end(),dfn[u])-col[w].begin();
	if (num < col[w].size() && col[w][num] <= dfn[v]) return 1;
	return 0;
}
int main() {
	scanf("%d%d",&n,&m);
	for (int i = 1;i <= n;i++) scanf("%d",&nod[i]);
	for (int i = 1,u,v;i < n;i++) {
		scanf("%d%d",&u,&v);
		edge[u].push_back(v);
		edge[v].push_back(u);
	}
	dfs1(1,0);
	dfs2(1,1);
	for (int i = 1;i <= n;i++) sort(col[i].begin(),col[i].end());
	for (int u,v,w;m--;) {
		scanf("%d%d%d",&u,&v,&w);
		printf("%d",query(u,v,w));
	}
	return 0;
}