题意
给定一棵树,每个节点有颜色c[i],定义子树u中出现次数最多的颜色为u的主要颜色(可以有多个),求每个节点的主要颜色数值和。
题解
树上启发合并
dus on tree
d
u
s
o
n
t
r
e
e
。
启发合并的意思就是说,将秩小的合并到秩大的上,有点并查集按秩合并的意思。而并查集的秩是和其下面的节点个数相关,树上也如是,对于树来说,还可以根据节点数目,划分轻重链,树上启发合并也从此而来。
由于轻重链划分时,重儿子不超过logn个,而树上启发合并的核心思想是,对每个待处理的子树问题,先找到的重儿子,统计答案的时候优先进入每一个点的所有轻儿子,之后再进入重儿子,目的是保留重儿子所在子树的信息。处理完当前点的所有儿子的子树之后开始处理自己。
那么就有:
- 先统计以当前点为根的子树不经过重儿子的所有点的影响 O(size[x]−size[hson[x]]) O ( s i z e [ x ] − s i z e [ h s o n [ x ] ] )
- 如果这个点是轻儿子则需要暴力删除这棵子树所带来的影响 O(size[x]) O ( s i z e [ x ] ) ,这也正是先进入轻儿子的原因,可以保留重儿子的信息。
对于每个儿子,被重复计算量为 O(到根节点的路径数目) O ( 到 根 节 点 的 路 径 数 目 ) ,可以确定他是 O(logn) O ( l o g n ) 级别的。
如此一来,就保证了时间复杂度。
下面说说我对保留重儿子信息的理解,可能需要结合部分代码,我把需要用到的代码放在下面
void update(int u,int f, ll num) {
cnt[color[u]] += num;
if(num>0 && cnt[color[u]] >= mxval){
if(cnt[color[u]] > mxval) nowans = 0, mxcolor = color[u], mxval = cnt[mxcolor];
nowans += color[u];
}
for(int i = head[u];i!=-1;i=e[i].nxt){
int v = e[i].to;
if(v != f && !visit[v]) update(v,u,num);
}
}
void dfsSecond(int u, int f, bool keep = false) {
for(int i = head[u];i!=-1;i=e[i].nxt){
int v = e[i].to;
if(v != f && v != son[u]) dfsSecond(v,u);
}
if(son[u]!=-1) dfsSecond(son[u],u,true), visit[son[u]] = true;
update(u,f,1); ans[u] = nowans;
if(son[u]!=-1) visit[son[u]] = false;
if(!keep) update(u,f,-1), nowans = mxcolor = mxval = 0;
}
其实就是上面两个函数。dfsSecond
就是求解问题的函数,而update
是更新答案的函数。
在dfsSecond
函数中,可以看到首先遍历所有儿子,并且这个儿子不是重儿子,就继续递归调用,所以最后的结果一定是递归到某个叶子节点,并且是一个轻儿子。
递归到叶子节点后,就跳出for循环,因为他没有儿子,由于他没有儿子,也就没有重儿子,所以son[u] = -1
,就会执行update
函数,而这个函数就会更新以当前叶子节点为根的子树的答案(其实就是他一个节点),计算完成后保存答案,由于是轻儿子,不会保存信息,所以再把答案清空。
接下来回溯,回溯到上一层,也就是他的父亲节点。他的父亲节点可能还有别的轻儿子,那么按照刚才说的思路执行。把它的所有轻儿子遍历后,去找他的重儿子。
下面重点来了,由于是重儿子,我们要保留信息,所以有dfsSecond(son[u],u,true)
,并且把它的重儿子置为访问过。递归进入重儿子后,首先也是遍历重儿子的轻儿子,统计结果。然后如果它还有重儿子,就继续先找重儿子,直到到达一个重儿子叶节点。
到达之后,由于他没有儿子,不会继续递归,转而进行update
函数,更新答案,此时注意,由于传递的keep = true
,所以不会清空答案(也就是以重儿子为根节点的子树信息,计算后不会清空,这点很重要)。之后回溯,回溯到的一定是他的父亲,如果他的父亲是重儿子,那么就进入update
统计答案,可以注意到,在update
中,只会对没有visit
的节点访问,而哪些节点visit
过了呢,显然是重儿子,换言之也就是指计算轻儿子的,而计算之后,回到dfsSecond
函数,由于本来节点是重儿子,所以传递的keep = true
,也就不会清空以此节点为根的所有信息,也就说明了以重儿子为根节点的子树信息,计算后不会清空。下次在用到重儿子的信息,直接统计就好了,保证复杂度。
代码
#include<bits/stdc++.h>
using namespace std;
typedef double db;
typedef long long ll;
typedef unsigned long long ull;
const int nmax = 1e6+7;
const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3f;
const ull p = 67;
const ull MOD = 1610612741;
int fa[nmax], son[nmax], sz[nmax], head[nmax], color[nmax];
ll ans[nmax], cnt[nmax], nowans, mxval, mxcolor;
bool visit[nmax];
int tot,n;
struct edge {int to, nxt;} e[nmax << 1];
void add(int u, int v) { e[tot].to = v, e[tot].nxt = head[u], head[u] = tot++;}
void dfsFirst(int u, int f) {
fa[u] = f, sz[u] = 1;
for(int i = head[u];i!= -1;i=e[i].nxt) {
int v = e[i].to;
if(v != f) {
dfsFirst(v, u);sz[u] += sz[v];
if (son[u] == -1 || sz[v] > sz[son[u]]) son[u] = v;
}
}
}
void update(int u,int f, ll num) {
cnt[color[u]] += num;
if(num>0 && cnt[color[u]] >= mxval){
if(cnt[color[u]] > mxval) nowans = 0, mxcolor = color[u], mxval = cnt[mxcolor];
nowans += color[u];
}
for(int i = head[u];i!=-1;i=e[i].nxt){
int v = e[i].to;
if(v != f && !visit[v]) update(v,u,num);
}
}
void dfsSecond(int u, int f, bool keep = false) {
for(int i = head[u];i!=-1;i=e[i].nxt){
int v = e[i].to;
if(v != f && v != son[u]) dfsSecond(v,u);
}
if(son[u]!=-1) dfsSecond(son[u],u,true), visit[son[u]] = true;
update(u,f,1); ans[u] = nowans;
if(son[u]!=-1) visit[son[u]] = false;
if(!keep) update(u,f,-1), nowans = mxcolor = mxval = 0;
}
int main(){
memset(head,-1,sizeof head);
memset(son,-1,sizeof son);
scanf("%d",&n);
for(int i = 1;i<=n;++i) scanf("%d", &color[i]);
int u,v;
for(int i = 1;i<n;++i) {
scanf("%d %d", &u, &v);
add(u,v);
add(v,u);
}
dfsFirst(1,0);
dfsSecond(1,0);
for(int i = 1;i<=n;++i) printf("%I64d ",ans[i]);
return 0;
}