洛谷U41492（树上启发式合并）

提交通道

洛谷日报

考虑非\(O(n^2)\)的预处理。一遍dfs时，check某颜色有没有的数组何时清空很尴尬：得到某树答案后如果不清，则影响接下来兄弟树的搜索；如果清了，父亲节点又难以收集答案。

解决方法：先让儿子们各顾各的家，算一遍各自的答案（假如能算），check清就清了吧。然后考虑人为优化，即重链求完后等一等！先别清！然后将轻链重新扫一遍，也不清check数组的。代码中的keep就控制是否要清。这样轻链扫两遍，重链扫一遍，就得到了儿子们和父亲的答案，随机数据下复杂度\(O(nlogn)\)。

const int maxn = 1e5 + 5;
int n, m, u, v, c[maxn];
int size[maxn], son[maxn], ans[maxn];
bool check[maxn];
vector<int> adj[maxn];

void dfs1(int cur, int fa) {
	size[cur] = 1;
	for (int i : adj[cur])
		if (i != fa) {
			dfs1(i, cur);
			size[cur] += size[i];
			if (size[i] > size[son[cur]])
				son[cur] = i;
		}
}

int dfs2(int cur, int fa, int keep) {
	if (keep) {
		for (int i : adj[cur])
			if (i != fa && i != son[cur])
				dfs2(i, cur, keep);
	}
	int res = 0;
	if (son[cur])	res += dfs2(son[cur], cur, keep);
	for (int i : adj[cur])
		if (i != fa && i != son[cur])
			res += dfs2(i, cur, 0);

	if (!check[c[cur]])	res++, check[c[cur]] = 1;
	if (keep) {
		ans[cur] = res;
		if (cur != son[fa])	mset(check, 0);
	}

	return res;
}

int main() {
	read(n);
	rep(i, 1, n - 1) {
		read(u), read(v);
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	rep(i, 1, n)	read(c[i]);

	dfs1(1, 0);
	dfs2(1, 0, 1);

	for (read(m); m--; ) {
		read(u);
		writeln(ans[u]);
	}
	return 0;
}