Codeforces 600E Lomsat gelral [dsu on tree(树上启发式合并)]

题意

给出一棵树，1为根节点，每个点都有一个颜色。求每个点所在子树内所有出现次数最多的颜色的和。

n<=100000

分析

新学习了一种姿势叫dsu on tree，大概意思就是树上启发式合并吧。

dsu on tree大概是用来解决这样一类问题：需要多次查询某棵子树内的某个值（必须要离线）。像这题就是需要查询每棵子树出现颜色次数最多的颜色的和。

怎么做呢？

首先我们把这棵树的dfs序求出来，顺便轻重链剖分一下，用一个全局数组统计当前答案。

每次到达一个节点，先访问其所有轻节点，每次暴力把轻节点子树的贡献删掉。然后访问重儿子，然后暴力把所有轻儿子的子树和自己的贡献加进去，最后查询当前点的询问即可。

复杂度的话，因为每个点到根的路径上，重链的数量不超过logn条，所以每个点最多被暴力加入logn次，于是总的复杂度就是O(nlogn).

伪代码

find the BigChild of each vertex
dfs(u, fa, keep)
dfs(LightChild, u, 0)
dfs(BigChild, u, 1), big[BigChild] = 1
update(u, fa, 1) //calculate the contribution of u's LightChild's SubTree
update the ans of u
big[BigChild] = 0
if keep == 0
update(u, fa, -1) //remove the contributino of u's SubTree

update(u, fa, val)
calculate u's information
update(v : (u, v) and !big[v], u, val)

具体到这题的话，只要开两个数组sum[i]表示出现了i次的颜色的和，tot[i]表示颜色i出现的次数，然后就变成模板题了。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
using namespace std;

typedef long long LL;

const int N=100005;

int n,mx
,mn
,sz,dfn
,tot
,col
,last
,size
,Mx,cnt;
struct edge{int to,next;}e[N*2];
LL sum
,ans
;

int read()
{
int x=0,f=1;char ch=getchar();
while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
return x*f;
}

void addedge(int u,int v)
{
e[++cnt].to=v;e[cnt].next=last[u];last[u]=cnt;
e[++cnt].to=u;e[cnt].next=last[v];last[v]=cnt;
}

void dfs(int x,int fa)
{
mn[x]=++sz;dfn[sz]=x;size[x]=1;
for (int i=last[x];i;i=e[i].next)
{
if (e[i].to==fa) continue;
dfs(e[i].to,x);
size[x]+=size[e[i].to];
}
mx[x]=sz;
}

void add(int x)
{
for (int i=mn[x];i<=mx[x];i++)
{
int y=col[dfn[i]];
sum[tot[y]]-=y;
sum[++tot[y]]+=y;
Mx=max(Mx,tot[y]);
}
}

void del(int x)
{
for (int i=mn[x];i<=mx[x];i++)
{
int y=col[dfn[i]];
sum[tot[y]]-=y;
sum[--tot[y]]+=y;
if (!sum[Mx]) Mx--;
}
}

void solve(int x,int fa)
{
int k=0;
for (int i=last[x];i;i=e[i].next)
if (e[i].to!=fa&&size[e[i].to]>size[k]) k=e[i].to;
for (int i=last[x];i;i=e[i].next)
if (e[i].to!=fa&&e[i].to!=k) solve(e[i].to,x),del(e[i].to);
if (k) solve(k,x);
for (int i=last[x];i;i=e[i].next)
if (e[i].to!=fa&&e[i].to!=k) add(e[i].to);
sum[tot[col[x]]]-=col[x];
sum[++tot[col[x]]]+=col[x];
Mx=max(Mx,tot[col[x]]);
ans[x]=sum[Mx];
}

int main()
{
n=read();
for (int i=1;i<=n;i++) col[i]=read(),sum[0]+=i;
for (int i=1;i<n;i++)
{
int x=read(),y=read();
addedge(x,y);
}
dfs(1,0);
solve(1,0);
for (int i=1;i<=n;i++) printf("%I64d ",ans[i]);
return 0;
}