树上启发式合并 洛谷U41492

对于子树类不修改的问题，比如说统计所有子树的颜色种类数，以下面这个样例为例子：

要统计以u为根的子树，应该暴力统计u的各个子树之和。流程应该是这样：假设规模为一个三个结点的二叉树，u为根，v1，v2为两个叶子节点，暴力统计以v1为根的节点，然后清空统计数组，再暴力统计以v2为根的节点，然后再清空数组，统计u为根的子树。我们可以发现，是因为两棵子树一定毫不相干我们才清空，但是v2与u不是毫不相干的，他们是包含关系的，所以我们可以保留v2的，那么在统计u的时候就可以不统计v2的了。更一般地，应该是保留u的最后一个儿子，节省统计u的时间，为了最大化节省的时间，我们就把重儿子放在最后一个（因为重儿子节点数最多，最耗时间），所以实际算法为了确保轻儿子一定在统计的前面，我们就一路轻儿子递归至叶子，再自底向根的回溯式统计。

统计代码：

void get_light_son(int u,int hson,int fa,int k)//统计轻儿子
{
	cnt[color[u]]+=k;
	if(k>0&&cnt[color[u]]==1) tot++;//一定要有k>0,因为可以从0->1也可以2->1,会多计数,只统计增加时的1
	if(k<0&&cnt[color[u]]==0) tot--;//同上
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if((k==1&&v==hson)||v==fa) continue;//如果v是fa或者是在增加的时候v是重儿子(被保留在cnt数组了，不需要再记录)，跳过
		get_light_son(v,hson,u,k);//递归的时候hson还是hson,因为我们只是不统计根的重儿子，即一开始传入的hson
	}
}

void dfs(int u,int fa)
{
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v!=fa&&v!=son[u]) dfs(v,u); //一路遍历轻儿子到叶子
	}
	if(son[u]) dfs(son[u],u);//重儿子放在最后统计
	get_light_son(u,son[u],fa,1);//计数这个节点为根的子树的所有节点，因为这个节点的重儿子被最后留下在cnt数组了，所以在内部特判重儿子，不记录
	ans[u]=tot;//以u为根的子树就是原先保留的重儿子加上get_light_son()的轻儿子记录，tot记录种数
	if(u!=son[fa]) get_light_son(u,son[u],fa,-1);//这个节点不是fa的重儿子，直接清空记录
}

其余就是链式前向星和统计树链剖分求重儿子的常规操作。

树上启发式合并的复杂度是的。证明待补。

完整代码：

#include<iostream>
#include<stack>
#include<queue>
#include<map>
#include<set>
#include<cctype>
#include<algorithm>
#include<string>
#include<vector>
#include<cmath>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<deque>
#define ll long long
#define ull unsigned long long
#define fr first
#define se second
#define endl '\n'
using namespace std;

const int inf=0x3fffffff;
const long long INF=0x3f3f3f3f3f3f;
const double eps=1e-8;
const double pi=2*acos(0.0);

template<typename T>
T father(T x){return x&1?x-1>>1:x>>1;}
template<typename T>
T ls(T x){return x<<1;}
template<typename T>
T rs(T x){return x<<1|1;}
template<typename T>
T abs_(T a){return a<0?-a:a;}
template<typename T>
T ceil_(T a,T b){return a%b?a/b+1:a/b;}
template<typename T>
T log_(T a,T b){return log(b)/log(a);}
template<typename T>
T gcd(T a,T b){return !b?a:gcd(b,a%b);}
template<typename T>
T lcm(T a,T b){return a/gcd(a,b)*b;}
template<typename T>
void exgcd(T a,T b,T& x,T& y)
{
	if(!b){x=1;y=0;return;}
	exgcd(b,a%b,x,y);
	T temp=x;x=y;
	y=temp-a/b*y;
}
template<typename T>
T solve_function(T a,T b,T c)
{
	if(a<0){a=-a;c=-c;}
	T x,y,base=gcd(a,b);
	exgcd(a,b,x,y);
	x*=c/base;T p=b/base;
	x=(x%p+p)%p;
	if(x<0) x+=abs_(p);
	return x;
}
template<typename T>
T reverse_element(T a,T k){return solve_function(a,k,(T)1);}
template<typename T>
T slow_multi(T a,T b,T p)
{
	T ret=0;
	while(b)
	{
		if(b&1) ret=(ret+a)%p;
		a=(a<<1)%p;
		b>>=1;
	}
	return ret;
}
template<typename T>
T qpow(T a,T b)
{
	T ret=1,base=a;
	while(b)
	{
		if(b&1) ret=ret*base;
		base=base*base;
		b>>=1;
	}
	return ret;
}

int read()
{
	int ret=0,base=1;
	char ch=getchar();
	while(!isdigit(ch))
	{
		if(ch=='-') base=-1;
		ch=getchar();
	}
	while(isdigit(ch))
	{
		ret=(ret<<1)+(ret<<3)+ch-48;
		ch=getchar();
	}
	return ret*base;
}

struct way
{
	int to,next;
}edge[200005];
int cnt2,head[100005];

void add(int u,int v)
{
	edge[++cnt2].to=v;
	edge[cnt2].next=head[u];
	head[u]=cnt2;
}

void add_(int u,int v){add(u,v);add(v,u);}

int n,color[100005],sum[100005],son[100005];
int maxCnt,tot,cnt[100005],ans[100005];

void dfs1(int u,int fa)
{
	sum[u]=1;
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v==fa) continue;
		dfs1(v,u);
		sum[u]+=sum[v];
		if(sum[v]>sum[son[u]]) son[u]=v;
	}
}

void get_light_son(int u,int hson,int fa,int k)//统计轻儿子
{
	cnt[color[u]]+=k;
	if(k>0&&cnt[color[u]]==1) tot++;
	if(k<0&&cnt[color[u]]==0) tot--;
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if((k==1&&v==hson)||v==fa) continue;
		get_light_son(v,hson,u,k);
	}
}

void dfs(int u,int fa)
{
	for(int i=head[u];i;i=edge[i].next)
	{
		int v=edge[i].to;
		if(v!=fa&&v!=son[u]) dfs(v,u); 
	}
	if(son[u]) dfs(son[u],u);
	get_light_son(u,son[u],fa,1);
	ans[u]=tot;
	if(u!=son[fa]) get_light_son(u,son[u],fa,-1);
}

int main()
{
	n=read();
	for(int i=1;i<n;i++) add_(read(),read());
	for(int i=1;i<=n;i++) color[i]=read();
	dfs1(1,0);dfs(1,0);
	for(int i=1;i<=n;i++) printf("%d\n",ans[i]);
 	return 0;
}