DSU on tree——令人惊叹的想法
DSU on tree
首先感谢LX dalao的讲解。
DSU on tree用于解决静态树上众数问题,比如说Codeforces 600E
题目大意
给你一棵树,每个节点有一种颜色,问你每个子树x的颜色数最多的那种颜色,如果颜色数相同,那么种类数相加。
输入样例1
4
1 2 3 4
1 2
2 3
2 4
输出样例1
10 9 3 4
输入样例2
15
1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
1 2
1 3
1 4
1 14
1 15
2 5
2 6
2 7
3 8
3 9
3 10
4 11
4 12
4 13
样例输出2
6 5 4 3 2 3 3 1 1 3 2 2 1 2 3
其实DSU on tree这个想法很简单,是一个神奇的优化算法,用树剖把笨蛋想法给优化到,我也不知道为什么时间复杂度是,但是他就是快。
我们先了解一下笨蛋的想法。
其实很简单,就是枚举这个点,然后这棵子树扫一遍得到答案,然后清空hsh数组。
这就是笨蛋,我们会发现它做了一些无用功,比如说最后一次清空,其实可以用于他的父节点,这样父节点就可以少算一个子节点。
我们想让尽量大的子树不擦除,那么就树剖剖出重儿子,重儿子不擦除就可以了!
这个想法很巧妙,先让我Orz一下发明人。
我们简单讲一下算法实现步骤:
- 剖出重儿子。
- DFS便利子树。先遍历轻儿子,然后再遍历重儿子。
- 如果当前根节点是重儿子,那么就跳过,否则就擦除。
下面解释一下代码:
变量:
int n,AnsMax,a[MAXN];
//n,a[]:读入
//AnsMax:当前子树最多的颜色
LL c,Ans[MAXN];
//c:当前的颜色
//Ans[]:存答案。
bool vis[MAXN];
//标记这个子树是否需要被遍历。
int hsh[MAXN],Siz[MAXN],Son[MAXN],Fa[MAXN];
//树剖
struct Edge{//邻接表
int tot,lnk[MAXN],son[MAXN<<1],nxt[MAXN<<1];
void Add(int x,int y){son[++tot]=y;nxt[tot]=lnk[x];lnk[x]=tot;}
}E;
树剖:
void First(int x,int f){
Fa[x]=f;Siz[x]=1;
for(int j=E.lnk[x],SizeMax=0;j;j=E.nxt[j])
if(E.son[j]!=f){
First(E.son[j],x);Siz[x]+=Siz[E.son[j]];
if(Siz[E.son[j]]>SizeMax) SizeMax=Siz[E.son[j]],Son[x]=E.son[j];
}
}
计算颜色:
void Cal(int x,int p){
hsh[a[x]]+=p;
//
if(p>0&&hsh[a[x]]==AnsMax) c+=a[x];
if(p>0&&hsh[a[x]]>AnsMax) AnsMax=hsh[a[x]],c=a[x];
//计算答案。
for(int j=E.lnk[x];j;j=E.nxt[j])
if(E.son[j]!=Fa[x]&&!vis[E.son[j]]) Cal(E.son[j],p);
//vis[E.son[j]]=1的话,说明hsh数组已经存了这棵子树
}
枚举子树:
void Second(int x,int T){//x子树根节点,T表示x是否是重儿子
for(int j=E.lnk[x];j;j=E.nxt[j])//遍历轻儿子
if(E.son[j]!=Fa[x]&&E.son[j]!=Son[x]) Second(E.son[j],0);
if(Son[x]) Second(Son[x],1),vis[Son[x]]=1;//遍历重儿子
Cal(x,1);Ans[x]=c;//计算整棵子树
if(Son[x]) vis[Son[x]]=0;
if(!T) Cal(x,-1),c=AnsMax=0;//如果不是重儿子,那么擦除
}
就这么简单
下面附上完整代码:
#include<cstdio>
#include<cctype>
#define MAXN 100005
#define LL __int64
using namespace std;
int n,AnsMax,a[MAXN];
LL c,Ans[MAXN];
bool vis[MAXN];
int hsh[MAXN],Siz[MAXN],Son[MAXN],Fa[MAXN];
struct Edge{
int tot,lnk[MAXN],son[MAXN<<1],nxt[MAXN<<1];
void Add(int x,int y){son[++tot]=y;nxt[tot]=lnk[x];lnk[x]=tot;}
}E;
int read(){
int ret=0;char ch=getchar();bool f=1;
for(;!isdigit(ch);ch=getchar()) f^=!(ch^'-');
for(; isdigit(ch);ch=getchar()) ret=(ret<<3)+(ret<<1)+ch-48;
return f?ret:-ret;
}
void First(int x,int f){
Fa[x]=f;Siz[x]=1;
for(int j=E.lnk[x],SizeMax=0;j;j=E.nxt[j])
if(E.son[j]!=f){
First(E.son[j],x);Siz[x]+=Siz[E.son[j]];
if(Siz[E.son[j]]>SizeMax) SizeMax=Siz[E.son[j]],Son[x]=E.son[j];
}
}
void Cal(int x,int p){
hsh[a[x]]+=p;
if(p>0&&hsh[a[x]]==AnsMax) c+=a[x];
if(p>0&&hsh[a[x]]>AnsMax) AnsMax=hsh[a[x]],c=a[x];
for(int j=E.lnk[x];j;j=E.nxt[j])
if(E.son[j]!=Fa[x]&&!vis[E.son[j]]) Cal(E.son[j],p);
}
void Second(int x,int T){
for(int j=E.lnk[x];j;j=E.nxt[j])
if(E.son[j]!=Fa[x]&&E.son[j]!=Son[x]) Second(E.son[j],0);
if(Son[x]) Second(Son[x],1),vis[Son[x]]=1;
Cal(x,1);Ans[x]=c;
if(Son[x]) vis[Son[x]]=0;
if(!T) Cal(x,-1),c=AnsMax=0;
}
int main(){
#ifndef ONLINE_JUDGE
freopen("prob.in","r",stdin);
freopen("prob.out","w",stdout);
#endif
n=read();
for(int i=1;i<=n;i++) a[i]=read();
for(int i=1;i<n;i++){int x=read(),y=read();E.Add(x,y);E.Add(y,x);}
First(1,0);Second(1,0);
for(int i=1;i<=n;i++) printf(i==n?"%I64d\n":"%I64d ",Ans[i]);
return 0;
}