dsu on tree(树上启发式合并)

/*
树上启发式合并
用于处理离线查询的子树问题
第一次先求出重儿子
第二次dfs的过程中,先计算轻儿子,贡献不保留,再计算重儿子,贡献保留
最后再加上轻儿子的贡献,计算出当前节点的值 
复杂度为O(nlogn) 
*/ 

/*
每个节点都有一个颜色编号 
计算一棵以1为根的树的以每个节点为根的子树中颜色最多的颜色编号和
这题中子树的贡献就在于cnt数组,用来计数颜色出现了多少次 
*/ 

#include <iostream>
#include <vector>
using namespace std;

const int maxn = 1e5 + 5;

typedef long long ll;

vector<int> g[maxn];
int c[maxn],cnt[maxn],son[maxn],siz[maxn];
ll ans[maxn];

void dfs1(int x,int f)   //找重儿子 
{
	siz[x] = 1;
	int maxsiz = 0;
	for (int i = 0; i < g[x].size(); i++)
	{
		int t = g[x][i];
		if( t == f ) continue;
		dfs1(t,x);
		siz[x] += siz[t];
		if( maxsiz < siz[t] )
		{
			son[x] = t;
			maxsiz = siz[t];
		}
	}
} 

int flag = 0;
ll sum = 0,maxc = 0;   //当前的答案和最大值 

void count(int x,int f,int val)  //将以x为根的子树的贡献加上val但不包括flag的贡献
{
	cnt[c[x]] += val;
	if( cnt[c[x]] > maxc )
	{
		maxc = cnt[c[x]];
		sum = c[x];
	}else if( cnt[c[x]] == maxc )  //两个颜色都为最大值都得算 
	{
		sum += c[x];
	}  	
	for (int i = 0;i < g[x].size(); i++)
	{
		int t = g[x][i];
		if( t == f || t == flag ) continue;
		count(t,x,val);
	}
}

void dfs2(int x,int f,int keep)  //当前节点为x,父亲节点为f,keep为false表示不保留贡献 
{
	for (int i = 0; i < g[x].size(); i++)  //处理轻儿子的子树,不保留贡献 
	{
		int t = g[x][i];
		if( t == f || t == son[x] ) continue;
		dfs2(t,x,false);
	}
	if( son[x] )
	{
		dfs2(son[x],x,true);   //处理重儿子的子树,保留贡献
		flag = son[x];    	   //标记重儿子,这样在算贡献时不会算上它 
	}
	count(x,f,1);    //将轻儿子和自身的贡献加上 
	ans[x] = sum;
	flag = 0;
	if( !keep ) 
	{
		count(x,f,-1);  	//keep为false不保留,以x为子树的贡献都加上-1,即消去了贡献 
		sum = maxc = 0;    //消去贡献后,sum和maxc都需要置为0 
	}//没消去贡献时,maxc与sum都是有用的 
} 

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	int n;
	cin >> n;
	for (int i = 1; i <= n; i++)
	{
		cin >> c[i];
	}
	for (int i = 1; i < n; i++) 
	{
		int x,y;
		cin >> x >> y;
		g[x].push_back(y);
		g[y].push_back(x);  
	}
	dfs1(1,0);
	dfs2(1,0,false);
	for (int i = 1; i <= n; i++)
	{
		cout << ans[i];
		if( i == n ) cout << '\n';
		else cout << ' ';
	}
	return 0;
}

证明复杂度：
对于一个节点需要再次被计算，当且仅当它的某个祖先节点不为其父亲节点的重儿子，假设这个节点被计算了logn次，由于有重儿子的存在，每次计算都会附带一个重儿子，节点数每次至少乘2，logn次后就达到了n，所以一个节点被计算次数不超过logn。