补题 dsu on tree CodeForces 600E Lomsat gelral

题目

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
The first line contains integer n (1 ≤ n ≤ 105) — 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
Print n integers — the sums of dominating colours for each vertex.
Examples
Input
4
1 2 3 4
1 2
2 3
2 4
Output
10 9 3 4
Input
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
Output
6 5 4 3 2 3 3 1 1 3 2 2 1 2 3

解题思路

题意就是树的每个点有一种颜色，当一种颜色在某一个顶点的子树中，数量最大时，称为占领了该子树，求其1到n所有的点，都是由那种颜色所占领

dsu on tree 的模板题，优先计算轻子树，然后再计算重子树，最后消除轻子树的影响

AC代码

/*
找出树中每一个节点的重儿子，统计答案的时候优先进入每一个点的所有轻儿子，之后再进入重儿子，目的是保留重儿子所在子树的信息。

处理完当前点的所有儿子的子树之后开始处理自己。

先统计以当前点为根的子树不经过重儿子的所有点的影响O(size[x]−size[hson[x]])
如果这个点是轻儿子则需要暴力删除这棵子树所带来的影响O(size[x])，这也正是先进入轻儿子的原因，可以保留重儿子的信息。
*/

#include
#include
#include
#include
using namespace std;
const int maxn=3*1e5+7;
typedef long long ll;
ll n,m,hson[maxn],sz[maxn],ans[maxn],c[maxn],cnt[maxn],sum,mx,s;
vector <ll> a[maxn];
void init(){
    ll x,y;
    cin>>n;
    for(ll i=1;i<=n;i++){
        cin>>c[i];
    }
    for(int i=1;i<n;i++){
        cin>>x>>y;
        a[x].push_back(y);
        a[y].push_back(x);//vector方式存储树
    }
}
void fhson(ll x,ll fa){//寻找重子树
    sz[x]=1;
    ll w=a[x].size();
    ll v;
    for(int i=0;i<w;i++){
        v=a[x][i];
        if(v==fa) continue;
        fhson(v,x);
        sz[x]+=sz[v];
        if(sz[v]>sz[hson[x]]) hson[x]=v;
    }
}
void calc(ll x,ll fa,ll v){//计算子树上的颜色的和
    cnt[c[x]]+=v;
    if(cnt[c[x]]>mx){
        sum=c[x];
        mx=cnt[c[x]];
    }
    else{
        if(cnt[c[x]]==mx) sum+=c[x];
    }
    ll w=a[x].size(),u;
    for(ll i=0;i<w;i++){
        u=a[x][i];
        if(u==fa||u==s) continue;//同样先计算轻子树
        calc(u,x,v);
    }
}
void dfs(ll x,ll fa,ll t){
    ll w=a[x].size();
    ll v;
    for(ll i=0;i<w;i++){
        v=a[x][i];
        if(v==fa||v==hson[x]) continue;//优先遍历其轻子树
        dfs(v,x,-1);
    }
    if(hson[x]) dfs(hson[x],x,1),s=hson[x];//记录该点的重子树
    calc(x,fa,1);//计算占领的颜色
    s=0;
    ans[x]=sum;
    if(t==-1){
        calc(x,fa,-1);//删除轻子树带来的影响，也就是把轻子树的点减去
        mx=sum=0;
    }
}


int main(){
    init();
    fhson(1,0);
    dfs(1,0,1);
    for(ll i=1;i<=n;i++){
        printf("%lld\n",ans[i]);
    }
    return 0;
}