Lomsat gelral CodeForces - 600E (tree heuristic merge)

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

Topic:
Give you a tree rooted in 1, each node has a color.

Ask you what color is the most colored subtree with roots at each node from 1 to n? If there are as many colors as there are, the answer should be their sum and.


Ideas:

Introduction to dsu on tree,

We know that the time complexity is n * n if we directly force a subtree rooted at each node, which obviously leads to tle.

We can optimize violence by exploiting the nature of tree's heavy son and light son, and the theoretical time complexity is O (nlogn).

We start with the root of the tree. For each node, we first deal with his younger son violently, maintain his answer, and empty the contribution of his younger son.

For the heavy son, we also deal with violence, but do not delete his contribution, because the heavy son node can contribute to its father node, that is, when we calculate the answer of the heavy son's father node, we do not need to sweep its heavy son, because it has already been processed.

At the same time, we can know the knowledge of tree chain partitioning. In this way, for each node, if he is a heavy son, he will only be visited once, and if he is a light son, he will visit logn at most. So the time complexity is O (nl ogn)

Recommend the blog that learns this knowledge: https://www.cnblogs.com/zwfymqz/p/9683124.html

See the code for details:

#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';
        }
    }
}


Tags: PHP iOS

Posted on Tue, 06 Aug 2019 02:14:24 -0700 by Virvo