图论-树上问题-树上启发式合并

树上启发式合并

前言

我们引入如下问题

给一棵根为1的树，每次询问子树颜色种类数。

想一个简单的方法，即定义 \(f[u][c]\) 表示以 \(u\) 为根，颜色 \(c\) 的出现次数。

不难发现，这种思路的时间与空间复杂度为 \(O(n^2)\) 所以本文章将介绍一种思路以解决这一类的问题。

介绍

启发式算法

启发式算法是基于人类的经验和直观感觉，对一些算法的优化。

举个例子，并查集，对于两个大小不一样的集合，我们将小的集合合并到大的集合中，优于将大的集合合并到小的集合中，因为并查集的特性，我们希望树的高度尽可能小，以查询其父节点

树上启发式合并

对于维护树上信息，需要将子节点维护的信息全部合并到父亲节点这一类型的问题时，由于空间和时间的限制，我们不能同时维护每一个节点的数据，所以对于上传数据的时候可以保留一部分信息来节省时间。不难想到，我们应保留数据更大即树的大小更大的部分，也就是树链剖分中常常提到的重链。此时就可以省去计算最重子树的数据，即时间复杂度可以到 \(O(nlogn)\) 空间可以达到 \(O(n)\)

具体实现：

1. 每次先进入轻节点，递归他们。
2. 进入重节点，递归他。
3. 统计轻节点，并添加信息。
4. 统计重节点，并添加信息。
5. 若当前点是轻子树，删除当前节点所有信息。

由于遍历轻节点时会删去部分数据，所以在部分题目下会影响答案，所以应先遍历轻节点，再遍历重节点。

代码

这里给出引入中问题的代码

#include <bits/stdc++.h>
#define int long long
using namespace std;
struct node{int to,next,w;}e[500010];int head[100010];int numEdge = 0;void add(int u, int v, int w){e[++numEdge].next = head[u];e[numEdge].to = v;e[numEdge].w = w;head[u] = numEdge;}
 
int col[100010],n = 0,m = 0, siz[100010], son[100010], tt = 0, dfn[100010], maxdfn[100010], nfd[100010], tot = 0, cnt[100010], mx = 0, ctc = 0;
 
void dfs(int u, int fa){
	siz[u] = 1;
	dfn[u] = ++ tt;
	nfd[tt] = u;
	for(int i = head[u]; i; i = e[i].next){
		int v = e[i].to;
		if(v == fa){
			continue;
		}
		dfs(v, u);
		siz[u] += siz[v];
		if(siz[v] > siz[son[u]]){
			son[u] = v;
		}
	}
	maxdfn[u] = tt;
}
void adds(int u, int fa){
	if(cnt[col[u]] == 0){
		ctc ++;
	}
	cnt[col[u]] ++;
}
void del(int u, int fa){
	cnt[col[u]] --;
	if(cnt[col[u]] == 0){
		ctc --;
	}
}
int ans[100010];
void dfs2(int u, int fa, bool flg){
	for(int i = head[u]; i; i = e[i].next){
		int v = e[i].to;
		if(v != fa && v != son[u]){
			dfs2(v, u, 0);
		}
	}
	if(son[u]){
		dfs2(son[u], u, 1);
	}
	for(int i = head[u]; i; i = e[i].next){
		int v = e[i].to;
		if(v != fa && v != son[u]){
			for(int i = dfn[v]; i <= maxdfn[v]; i++){
				adds(nfd[i], u);
			}
		}
	}
	adds(u, fa);
	ans[u] = ctc;
	if(!flg){
		for(int i = dfn[u]; i <= maxdfn[u]; i++){
			del(nfd[i], u);
		}
	}
}
signed main(){
	cin >> n;
	
	for(int i = 1; i < n; i++){
		int u = 0;
		int v = 0;
		cin >> u >> v;
		add(u, v, 1);
		add(v, u, 1);
	}
	for(int i = 1; i <= n; i++){
		cin >> col[i];
	}
	dfs(1, 0);
	dfs2(1, 0, 0);
	int m = 0; cin >> m;
	while(m --){
		int op = 0;
		cin >> op;
		cout << ans[op] << "\n";
	}
 
	return 0;
}