洛谷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;

}

巴特西

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

最新文章

热门文章