Lomsat gelral CodeForces - 600E （树上启发式合并）

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 ci (1 ≤ ci ≤ n), ci — the colour of the i-th vertex.

Each of the next n - 1 lines contains two integers xj, yj (1 ≤ xj, yj ≤ 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为根的树，每一个节点有一个颜色。

询问你对于从1到n每一个节点为根的子树中，颜色最多的是哪个颜色？如果有多个颜色数量一样多，答案应该是他们的sum和。

思路：

dsu on tree 的入门题，

我们知道如果直接暴力求对于每一个节点为根的子树话，时间复杂度是 n * n的，显然会tle，的

我们可以利用树的重儿子和轻儿子的性质来优化暴力，而理论的时间复杂度是 O（nlogn）

我们从树根开始dfs，对于每一个节点，我们先暴力处理他的轻儿子，维护出清儿子的答案，同时清空轻儿子的贡献。

而对于重儿子，我们同样暴力处理，但是不删除他的贡献，因为重儿子节点可以对它的父节点有贡献，即我们在算重儿子的父节点的答案时，就不需要去扫它的重儿子了，因为已经处理过了。

同时，树链剖分的知识我们可以知道，这样处理的话，对于每一个节点，如果他是重儿子，他只会被访问1次，如果是轻儿子，最多访问logn次。所以时间复杂度是 O（nl ogn ）

推荐学习本知识点的博客：https://www.cnblogs.com/zwfymqz/p/9683124.html

细节见代码：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <vector>
#include <iomanip>
#define ALL(x) (x).begin(), (x).end()
#define sz(a) int(a.size())
#define all(a) a.begin(), a.end()
#define rep(i,x,n) for(int i=x;i<n;i++)
#define repd(i,x,n) for(int i=x;i<=n;i++)
#define pii pair<int,int>
#define pll pair<long long ,long long>
#define gbtb ios::sync_with_stdio(false),cin.tie(0),cout.tie(0)
#define MS0(X) memset((X), 0, sizeof((X)))
#define MSC0(X) memset((X), '\0', sizeof((X)))
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define eps 1e-6
#define gg(x) getInt(&x)
#define chu(x) cout<<"["<<#x<<" "<<(x)<<"]"<<endl
using namespace std;
typedef long long ll;
ll gcd(ll a, ll b) {return b ? gcd(b, a % b) : a;}
ll lcm(ll a, ll b) {return a / gcd(a, b) * b;}
ll powmod(ll a, ll b, ll MOD) {ll ans = 1; while (b) {if (b % 2)ans = ans * a % MOD; a = a * a % MOD; b /= 2;} return ans;}
inline void getInt(int* p);
const int maxn = 100010;
const int inf = 0x3f3f3f3f;
/*** TEMPLATE CODE * * STARTS HERE ***/

int n;
std::vector<int> son[maxn];
int wson[maxn];
int SZ[maxn];
int a[maxn];
void dfs1(int x,int pre)
{
    SZ[x]=1;
    int maxson=-1;
    for(auto y:son[x])
    {
        if(y!=pre)
        {
            dfs1(y,x);
            SZ[x]+=SZ[y];
            if(maxson<SZ[y])
            {
                maxson=SZ[y];
                wson[x]=y;
            }
        }
    }
}
ll ans[maxn];
ll sum;
int isson;
int m;
ll cnt[maxn];
void add(int x,int pre,int val)
{
    cnt[a[x]]+=val;
    if(cnt[a[x]]>m)
    {
        m=cnt[a[x]];
        sum=a[x];
    }else if(cnt[a[x]]==m)
    {
        sum+=a[x];
    }
    for(auto y:son[x])
    {
        if(y==pre||y==isson)
            continue;
        add(y,x,val);
    }
}
void dfs2(int x,int pre,int op)
{
    for(auto y:son[x])
    {
        if(y==pre||y==wson[x])
        {
            continue;
        }
        dfs2(y,x,0);
    }
    if(wson[x])
    {
        dfs2(wson[x],x,1);
        isson=wson[x];
    }
    add(x,pre,1);
    isson=0;
    ans[x]=sum;
    if(op==0)
    {
        add(x,pre,-1);
        sum=0;
        m=0;
    }
}
int main()
{
    //freopen("D:\\code\\text\\input.txt","r",stdin);
    //freopen("D:\\code\\text\\output.txt","w",stdout);
    gg(n);
    repd(i,1,n)
    {
        gg(a[i]);
    }
    int u,v;
    repd(i,2,n)
    {
        gg(u);gg(v);
        son[u].pb(v);
        son[v].pb(u);
    }
    dfs1(1,0);
    dfs2(1,0,0);

    repd(i,1,n)
    {
        printf("%lld ",ans[i] );
    }
    printf("\n");


    return 0;
}

inline void getInt(int* p) {
    char ch;
    do {
        ch = getchar();
    } while (ch == ' ' || ch == '\n');
    if (ch == '-') {
        *p = -(getchar() - '0');
        while ((ch = getchar()) >= '0' && ch <= '9') {
            *p = *p * 10 - ch + '0';
        }
    }
    else {
        *p = ch - '0';
        while ((ch = getchar()) >= '0' && ch <= '9') {
            *p = *p * 10 + ch - '0';
        }
    }
}