Heuristic merging simple problem on tree

Lomsat gelral( Process in notes )

Topic translation

• There is a tree $n$ The number of nodes is $1$ Node number is the root of Rooting tree .
• Each node has a color , Colors are represented by numbers , $i$ The color number of node No. is $c_i$.
• If a color is in $x$ The subtree with root appears the most times , Call it in $x$ In a subtree with roots Dominance . obviously , Multiple colors may dominate in the same subtree .
• Your task is for every $i\in [1,n]$, Find out to $i$ In the subtree of the root , The number and number of the dominant colors .
• $n\le 10^5,c_i\le n$

Title Description

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$ .

Input format

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.

Output format

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

Examples #1

The sample input #1

4
1 2 3 4
1 2
2 3
2 4

Sample output #1

10 9 3 4

Examples #2

The sample input #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

Sample output #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){
    // Chain forward star edge 
	ed[cnt]=stuu{
    v,head[u]};
	head[u]=cnt++;
}
void dfs1(int u,int f){
    //u Is the current node ,f Is the parent node of the current node ; initialization 1
    siz[u]=1;
    int maxsize=-1;// It is a temporary variable to judge whether it is a heavy son 
    for(int i=head[u];i!=-1;i=ed[i].nt){
    // Traverse all the sons , Keep updating and find your son 
		int v=ed[i].to;
        if(v==f) continue;// My father must have skipped it directly 
        dfs1(v,u);// Depth traversal , The current node becomes the parent node , The found son becomes the current node and continues to traverse 
        siz[u]+=siz[v];// After traversal , Let the size of the current node plus the size of the son 
        if(siz[v]>maxsize){
    // If the size of the son is larger than the temporary variable 
            maxsize=siz[v];// Just assign it to a temporary variable 
            son[u]=v;// Change the heavy son of the current node 
        }
    }
}
void count(int u,int f,int val){
    // Count the contributions of a node and all its young sons 
    cntt[col[u]]+=val;//val Positive or negative can control whether to increase or delete contributions 
    if(cntt[col[u]]>kmax){
    // look for max
        kmax=cntt[col[u]];
        sum=col[u];// Find out u In the subtree of the root , The number and number of the dominant colors 
    }
    else if(cntt[col[u]]==kmax) sum+=col[u];// If the number of two colors is the same, it should be counted 
    for(int i=head[u];i!=-1;i=ed[i].nt){
    // Exclude the marked heavy son , Statistics of other son subtree information 
        int v=ed[i].to;
        if(v==f||v==flag) continue;
        count(v,u,val);
    }
}
void dfs2(int u,int f,int keep){
    
    // Traverse all light sons and their subtrees to calculate their answer deletion contribution 
    for(int i=head[u];i!=-1;i=ed[i].nt){
    // Traverse all the light sons 
        int v=ed[i].to;
        if(v==f||v==son[u]) continue;// If it is a parent node or a son, skip 
        dfs2(v,u,false);
    }
    // Deal with the heavy son and his subtree to calculate his answer without deleting the contribution 
    if(son[u]){
    
        dfs2(son[u],u,true);
        flag=son[u];// Mark heavy son , It is convenient to skip 
    }
    count(u,f,1);// Violence statistics u And all the contributions of the light son are merged into the heavy son information just calculated 
    flag=0;
    ans[u]=sum;
    if(!keep) count(u,f,-1),sum=kmax=0;// Delete the contributions that need to be deleted   And clear it for next use 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]);
}