【CF600E】Lomsat gelral

题目

题目链接：https://codeforces.com/problemset/problem/600/E
有一棵 \(n\) 个结点的以 \(1\) 号结点为根的有根树。
每个结点都有一个颜色，颜色是以编号表示的， \(i\) 号结点的颜色编号为 \(c_i\)。
如果一种颜色在以 \(x\) 为根的子树内出现次数最多，称其在以 \(x\) 为根的子树中占主导地位。显然，同一子树中可能有多种颜色占主导地位。
你的任务是对于每一个 \(i\in[1,n]\)，求出以 \(i\) 为根的子树中，占主导地位的颜色的编号和。
\(n\le 10^5,c_i\le n\)。

思路

一眼线段树合并，但是被安排用 dsu on tree 做 /kk。
设 \(cnt[x]\) 表示颜色 \(x\) 出现的次数。那么我们对于一个点 \(x\)，先求出其轻儿子的答案并将 \(cnt\) 清空，然后求重儿子的答案并不清空。接下来再便利轻子树的每一个节点把贡献加上即可。
显然复杂度瓶颈在于每一个点最多做多少次贡献。一个点被计算次数取决于从它到根的轻边条数。再树剖中我们都知道，重链是 \(O(\log n)\) 级别的，每两个重链之间夹着恰好一个轻边，所以每一个点贡献次数是 \(O(\log n)\) 的。
时间复杂度 \(O(n\log n)\)。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=100010;
int n,tot,head[N],cnt[N],col[N],size[N],son[N];
ll ans[N],maxn,sum;

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	e[++tot].to=to;
	e[tot].next=head[from];
	head[from]=tot;
}

void dfs1(int x,int fa)
{
	size[x]=1;
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa)
		{
			dfs1(v,x);
			size[x]+=size[v];
			if (size[v]>size[son[x]]) son[x]=v;
		}
	}
}

void check(int x)
{
	if (cnt[col[x]]>maxn)
		maxn=cnt[col[x]],sum=col[x];
	else if (cnt[col[x]]==maxn)
		sum+=col[x];
}

void dfs3(int x,int fa,int v)
{
	cnt[col[x]]+=v;
	check(x);
	for (int i=head[x];~i;i=e[i].next)
		if (e[i].to!=fa) dfs3(e[i].to,x,v);
}

void dfs2(int x,int fa,bool clr)
{
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa && v!=son[x]) dfs2(v,x,1);
	}
	sum=maxn=0;
	if (son[x]) dfs2(son[x],x,0);
	cnt[col[x]]++;
	check(x);
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa && v!=son[x]) dfs3(v,x,1);
	}
	ans[x]=sum;
	if (clr)
	{
		cnt[col[x]]--;
		for (int i=head[x];~i;i=e[i].next)
		{
			int v=e[i].to;
			if (v!=fa) dfs3(v,x,-1);
		}
		sum=maxn=0;
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d",&n);
	for (int i=1;i<=n;i++)
		scanf("%d",&col[i]);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	dfs1(1,0); dfs2(1,0,0);
	for (int i=1;i<=n;i++)
		printf("%lld ",ans[i]);
	return 0;
}