[CF600] E. Lomsat gelral ( 经典树上启发式合并 dus on tree )

题目链接：E. Lomsat gelral

题解

经典树上启发式合并的题目，按照正常暴力求时间复杂度为$O(n^2)$，很明显是有问题的。那么树上启发式合并是如何做的呢。

先说一些前导知识

1. 重儿子/轻儿子
   求出每棵子树的大小，节点数最多的一个子树的根节点就是这个节点父节点的重儿子。
   轻儿子即是节点个数最少的子树的根节点。
2. 重边/轻边
   重边即是连接重儿子与其父亲的边。
   轻边即是连接轻儿子与其父亲的边。

设轻儿子v、父亲节点f，由定义可推2*size(v)<size(f)，易得从一个节点到根节点的轻边数量不超过logn 条。现在我们考虑，如果先遍历所有的轻儿子子树，最后再遍历重儿子子树，那么时间复杂度是多少？

由于一个节点可能有多个轻儿子，所以每当我们暴力统计完一个轻儿子再去统计另一个轻儿子时，就得把之前那个轻儿子子树统计的信息全部删除，同样的道理如果这个轻儿子的父节点也是轻儿子节点时，父节点被统计时会再次遍历到该节点，由于一个节点到根节点的轻边数量不超过logn 条，可得每个轻儿子节点最多被统计logn次。
由于重儿子是最后被统计的，所以没有必要删除重儿子子树的节点信息，即使重儿子的父节点是轻儿子节点，重儿子被访问的次数也不会超过logn次，

综上所述，时间复杂度为O(nlogn)，其思想就是保留重儿子的结果，暴力迭代轻儿子。

为什么这个叫树上启发式合并呢？其实，求解的过程可以看作是轻儿子不断往重儿子合并的过程，这和常规的启发式合并是一致的。从这个角度来看，复杂度也是 O(nlogn)。

代码

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cstring>
#include<bitset>
#include<cassert>
#include<cctype>
#include<cmath>
#include<cstdlib>
#include<ctime>
#include<deque>
#include<iomanip>
#include<list>
#include<map>
#include<queue>
#include<set>
#include<stack>
#include<vector>
#include<unordered_set>
#include<unordered_map>
using namespace std;
//extern "C"{void *__dso_handle=0;}
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;
#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define pii pair<int,int>
#define lowbit(x) x&-x

const double PI=acos(-1.0);
const double eps=1e-6;
const ll mod=1e9+7;
const int inf=0x3f3f3f3f;
const int maxn=1e5+10;
const int maxm=1e6+10;
#define ios ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

int col[maxn],son[maxn],siz[maxn],mx,cnt[maxn],Son;
ll sum,ans[maxn];
vector<int> g[maxn];
void dfs_son(int x,int fa)
{
	siz[x] = 1;
	for(int i=0;i<g[x].size();i++)
	{
		int v=g[x][i];
		if(v == fa) continue;
		dfs_son(v, x);
		siz[x] += siz[v];
		if(siz[v]>siz[son[x]]) son[x]=v;
	}
	return ;
}

void update(int x,int fa,int val)
{
	cnt[col[x]]+=val;
	if(cnt[col[x]]>mx) mx=cnt[col[x]],sum=col[x];
	else if(cnt[col[x]]==mx) sum+=(ll)col[x];
	
	for(int i=0;i<g[x].size();i++)
	{
		int v=g[x][i];
		if(v==fa || v==Son) continue;
		update(v, x, val);
	}
}

void dfs2(int x,int fa,int op)
{
	for(int i=0;i<g[x].size();i++)
	{
		int v=g[x][i];
		if(v==fa || v==son[x]) continue;
		dfs2(v, x, 0);
	}
	if(son[x]) dfs2(son[x], x, 1),Son=son[x];
	update(x,fa,1); Son=0;
	ans[x]=sum;
	if(!op) update(x, fa, -1),mx=sum=0;
}

int main()
{
	int n; scanf("%d",&n);
	for(int i=1;i<=n;i++) scanf("%d",&col[i]);
	for(int i=1;i<n;i++)
	{
		int u,v; scanf("%d%d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs_son(1, 0);
	dfs2(1, 0, 1);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
}