1. 程式人生 > >Codeforces 600E. Lomsat gelral(樹上啟發式合併)

Codeforces 600E. Lomsat gelral(樹上啟發式合併)

題意:

n個點的有根樹,以1為根,每個點有一種顏色。我們稱一種顏色佔領了一個子樹當且僅當沒有其他顏色在這個子樹中出現得比它多。求佔領每個子樹的所有顏色之和。

題解:

#include<bits/stdc++.h>
using namespace std;
#define N 100005
#define LL long long
int n, Max, cnt[N], sz[N], col[N], big[N];
LL sum, ans[N];
vector<int> g[N];
void add(int u, int fa, int x)
{
    cnt[col[u]]+=x;
    if(cnt[col[u]]>Max) sum=col[u], Max=cnt[col[u]];
    else if(cnt[col[u]]==Max) sum+=col[u];
    for(auto v:g[u])
        if(v!=fa&&!big[v])
            add(v, u, x);
}
void dfs(int u, int fa, bool keep)
{
    int mx=-1, bigchild=-1;
    for(auto v:g[u])
        if(v!=fa&&sz[v]>mx)
            mx=sz[v], bigchild=v;
    for(auto v:g[u])
        if(v!=fa&&v!=bigchild)
            dfs(v, u, 0);
    if(bigchild!=-1)
        dfs(bigchild, u, 1), big[bigchild]=1;
    add(u, fa, 1);
    if(bigchild!=-1)
        big[bigchild]=0;
    ans[u]=sum;
    if(keep==0)
        add(u, fa, -1), Max=sum=0;
}

void getsz(int u, int fa)
{
    sz[u]=1;
    for(auto v:g[u])
        if(v!=fa)
        {
            getsz(v, u);
            sz[u]+=sz[v];
        }
}
void solve()
{
    cin>>n;
    int u, v;
    for(int i=1; i<=n; i++) cin>>col[i];
    for(int i=1; i<n; i++) cin>>u>>v, g[u].push_back(v), g[v].push_back(u);
    getsz(1, 0);
    dfs(1, 0, 1);
    for(int i=1; i<=n; i++) cout << ans[i] << " ";
}
int main()
{
    solve();
    return 0;
}