dsu on tree 树上启发式合并

模板题
由于共用一个桶,所以每次求完子树答案后需要把这个子树的点从桶中删除,很显然这就过于暴力了;
我们可以注意到如果我们算完一个子树后马上算他的父亲,那么这个子树的点就是没必要删掉的;
那么我们可以贪心的把size最大的子树放到最后来算,并且算完后不删除他的点;
然后就做完了...
看起来还是很暴力?但是可以证明这是一个nlogn的做法 不过我不会证
代码:

#include<bits/stdc++.h>
using namespace std;
const int N = 1e5 + 5;
vector<int> mp[N];
int n, m, tot, cnt, dfn[N], c[N], ans[N], son[N], in[N], out[N], flag[N], size[N], tong[N];
void add(int x){
	tong[x]++;
	if(tong[x] == 1) tot++;
}
void sub(int x){
	tong[x]--;
	if(tong[x] == 0) tot--;
}
void dfs(int u, int fa){ //跑出dfs序、该子树dfs序起点和终点,以及最重儿子
	in[u] = dfn[u] = ++cnt;
	flag[cnt] = u;
	size[u]++;
	for(auto v : mp[u]){
		if(v == fa) continue;
		dfs(v, u);
		size[u] += size[v];
		if(size[v] > size[son[u]]) son[u] = v;
	}
	out[u] = cnt;
}
void dst(int u, int fa, bool keep){ //keep表示在不在桶中删它的点
	for(auto v : mp[u]){ //先算轻儿子的答案
		if(v == fa || v == son[u]) continue;
		dst(v, u, 0);
	}
	if(son[u]) dst(son[u], u, 1); //后算重儿子
	for(int i = in[u]; i <= out[u]; i++){
		if(i == dfn[son[u]]){ //重儿子的点已经在桶中,跳过整个子树
			i = out[son[u]];
			continue;
		}
		add(c[flag[i]]);
	}
	ans[u] = tot;
	if(!keep){
		for(int i = in[u]; i <= out[u]; i++) sub(c[flag[i]]);
	}
}
int main(){
	cin >> n;
	for(int i = 1; i < n; i++){
		int u, v;
		cin >> u >> v;
		mp[u].push_back(v);
		mp[v].push_back(u);
	}
	dfs(1, 0);
	for(int i = 1; i <= n; i++) cin >> c[i];
	dst(1, 0, 1);
	cin >> m;
	for(int i = 1; i <= m; i++){
		int x;
		cin >> x;
		cout << ans[x] << '\n'; 
	}
	return 0;
}
posted @ 2026-01-10 16:25  Turkey_VII  阅读(13)  评论(1)    收藏  举报