树上启发式合并简单题

Lomsat gelral(过程在注释)

题面翻译

• 有一棵 $n$ 个结点的以 $1$ 号结点为根的有根树。
• 每个结点都有一个颜色，颜色是以编号表示的， $i$ 号结点的颜色编号为 $c_i$。
• 如果一种颜色在以 $x$ 为根的子树内出现次数最多，称其在以 $x$ 为根的子树中占主导地位。显然，同一子树中可能有多种颜色占主导地位。
• 你的任务是对于每一个 $i\in [1,n]$，求出以 $i$ 为根的子树中，占主导地位的颜色的编号和。
• $n\le 10^5,c_i\le n$

题目描述

You are given a rooted tree with root in vertex $1$ .

Each vertex is coloured in some colour.

Let’s call colour $c$ dominating in the subtree of vertex $v$ if there are no other colours that appear in the subtree of vertex $v$ more times than colour $c$ .

So it’s possible that two or more colours will be dominating in the subtree of some vertex.

The subtree of vertex $v$ is the vertex $v$ and all other vertices that contains vertex $v$ in each path to the root.

For each vertex $v$ find the sum of all dominating colours in the subtree of vertex $v$ .

输入格式

The first line contains integer $n$ ( $1<=n<=10^5$ ) — the number of vertices in the tree.

The second line contains $n$ integers $c_i$ ( $1<=c_i<=n$ ), $c_i$ — the colour of the $i$ -th vertex.

Each of the next $n-1$ lines contains two integers $x_j,y_j$ ( $1<=x_j,y_j<=n$ ) — the edge of the tree.

The first vertex is the root of the tree.

输出格式

Print $ n $ integers — the sums of dominating colours for each vertex.

样例 #1

样例输入 #1

4
1 2 3 4
1 2
2 3
2 4

样例输出 #1

10 9 3 4

样例 #2

样例输入 #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

#include<bits/stdc++.h>
using namespace std;
const int N=2e5+10;
typedef long long ll;
int head[N],cnt;
int siz[N],son[N];
int col[N],cntt[N];
ll ans[N],sum;
int flag,kmax;
struct stuu{
    
	int to,nt;
}ed[N*2];
void add(int u,int v){
    //链式前向星加边
	ed[cnt]=stuu{
    v,head[u]};
	head[u]=cnt++;
}
void dfs1(int u,int f){
    //u为当前节点,f为当前节点的父节点；初始化1
    siz[u]=1;
    int maxsize=-1;//判断是不是重儿子的临时变量
    for(int i=head[u];i!=-1;i=ed[i].nt){
    //遍历所有儿子，不断更新同时找到重儿子
		int v=ed[i].to;
        if(v==f) continue;//是父亲肯定直接跳过
        dfs1(v,u);//深度遍历,当前节点变为父节点，找到的儿子变为当前节点继续遍历下去
        siz[u]+=siz[v];//遍历完成后，让当前节点的大小加上儿子的大小
        if(siz[v]>maxsize){
    //如果儿子的大小大于临时变量
            maxsize=siz[v];//就赋给临时变量
            son[u]=v;//更改当前节点的重儿子
        }
    }
}
void count(int u,int f,int val){
    //统计某节点及其所有轻儿子的贡献
    cntt[col[u]]+=val;//val为正为负可以控制是增加贡献还是删除贡献
    if(cntt[col[u]]>kmax){
    //找max
        kmax=cntt[col[u]];
        sum=col[u];//找以u为根的子树中，占主导地位的颜色的编号和
    }
    else if(cntt[col[u]]==kmax) sum+=col[u];//如果两个颜色数量相同那都要算
    for(int i=head[u];i!=-1;i=ed[i].nt){
    //排除被标记的重儿子，统计其它儿子子树信息
        int v=ed[i].to;
        if(v==f||v==flag) continue;
        count(v,u,val);
    }
}
void dfs2(int u,int f,int keep){
    
    //遍历所有轻儿子及其子树算其答案删贡献
    for(int i=head[u];i!=-1;i=ed[i].nt){
    //遍历所有轻儿子
        int v=ed[i].to;
        if(v==f||v==son[u]) continue;//是父节点或重儿子就跳过
        dfs2(v,u,false);
    }
    //处理重儿子及其子树算其答案不删贡献
    if(son[u]){
    
        dfs2(son[u],u,true);
        flag=son[u];//标记重儿子，方便统计贡献时跳过
    }
    count(u,f,1);//暴力统计u及其所有轻儿子的贡献合并到刚算出的重儿子信息里
    flag=0;
    ans[u]=sum;
    if(!keep) count(u,f,-1),sum=kmax=0;//把需要删贡献的删掉 并为下次使用清0
}
int main(){
    
    memset(head,-1,sizeof(head));
    int n;
    cin>>n;
    for(int i=1;i<=n;i++) scanf("%d",&col[i]);
    int u,v;
    for(int i=1;i<n;i++){
    
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    dfs1(1,0),dfs2(1,0,0);
    for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
}