CF EDU - E. Lomsat gelral 树上启发式合并

学习：http://codeforces.com/blog/entry/44351

E. Lomsat gelral 

题意：

给定一个以1为根节点的树，每个节点都有一个颜色，问每个节点的子树中，颜色最多的是哪几种颜色，输出这些颜色的值得和。

 

思路：

树上启发式合并的模板题，具体来说，先对树进行树链剖分，分出一个重链和轻边，dfs时，把每条轻儿子暴力加到根节点中，每次加的时候，用这样的技巧

if(csz < cnt[col[v]]) sum = col[v] , csz = cnt[col[v]];
else if(csz==cnt[col[v]]) sum += col[v];

就可以使得sum中保存的是这个节点的子树中数量相同且最多的颜色和。

 

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cstring>
#include <string>
#include <vector>
#include <map>
#include <set>
#include <queue>
#include <list>
#include <cstdlib>
#include <iterator>
#include <cmath>
#include <iomanip>
#include <bitset>
#include <cctype>
using namespace std;
//#pragma comment(linker, "/STACK:102400000,102400000")  //c++
#define lson (l , mid , rt << 1)
#define rson (mid + 1 , r , rt << 1 | 1)
#define debug(x) cerr << #x << " = " << x << "\n";
#define pb push_back
#define pq priority_queue

typedef long long ll;
typedef unsigned long long ull;

typedef pair<ll ,ll > pll;
typedef pair<int ,int > pii;

//priority_queue<int> q;//这是一个大根堆q
//priority_queue<int,vector<int>,greater<int> >q;//这是一个小根堆q
#define fi first
#define se second
#define endl '\n'

#define OKC ios::sync_with_stdio(false);cin.tie(0)
#define FT(A,B,C) for(int A=B;A <= C;++A)  //用来压行
#define REP(i , j , k)  for(int i = j ; i <  k ; ++i)
//priority_queue<int ,vector<int>, greater<int> >que;

const ll mos = 0x7FFFFFFF;  //2147483647
const ll nmos = 0x80000000;  //-2147483648
const int inf = 0x3f3f3f3f;
const ll inff = 0x3f3f3f3f3f3f3f3f; //18

template<typename T>
inline T read(T&x){
    x=0;int f=0;char ch=getchar();
    while (ch<'0'||ch>'9') f|=(ch=='-'),ch=getchar();
    while (ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    return x=f?-x:x;
}
// #define _DEBUG;         //*//
#ifdef _DEBUG
freopen("input", "r", stdin);
// freopen("output.txt", "w", stdout);
#endif
/*-----------------show time----------------*/
                const int maxn = 1e5+9;
                vector<int>g[maxn];
                vector<int>v[maxn];
                int cnt[maxn],col[maxn];
                ll ans[maxn];
                bool big[maxn];
                int n;
                int sz[maxn];

                void dfs1(int v,int fa){
                    sz[v] = 1;
                    for(int i=0; i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa){
                            dfs1(u,v);
                            sz[v]+=sz[u];
                        }
                    }
                }
                ll sum = 0,csz = 0;
                // priority_queue<int>que;
                void add(int v, int fa,int x){

                    cnt[col[v]] += x;
                    // if(v==1)cout<<"$$"<<cnt[col[v]]<<endl;
                    if(csz < cnt[col[v]]) sum = col[v] , csz = cnt[col[v]];
                    else if(csz==cnt[col[v]]) sum += col[v];

                    for(int i=0;i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa && big[u]==0){
                            add(u,v,x);
                        }
                    }
                }

                void dfs(int v,int fa,int keep){
                    int mx = -1,bigChild = -1;

                    for(int i=0; i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa && sz[u] > mx)
                            mx = sz[u], bigChild = u;
                    }

                    for(int i=0; i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa && u!=bigChild){
                            dfs(u,v,0);////////////
                        }
                    }

                    if(bigChild!=-1){
                        dfs(bigChild,v,1),big[bigChild] = 1;/////////
                    }

                    add(v,fa,1);            ////////
                    ans[v] = sum;
                    /*
                    debug(v);
                    for(int i=1; i<=n; i++){
                        cout<<cnt[i]<<" ";
                    }
                    cout<<endl;
                    */
                    if(bigChild!=-1) big[bigChild] = 0;
                    if(keep==0){
                        add(v,fa,-1);
                        sum = 0;csz = 0;
                    }
                    /*
                    debug(v);
                    for(int i=1; i<=n; i++){
                        cout<<cnt[i]<<" ";
                    }
                    cout<<endl;
                    */
                }

int main(){
                scanf("%d", &n);
                for(int i=1; i<=n; i++)scanf("%d", &col[i]);
                for(int i=1; i<n; i++){
                    int u,v;
                    scanf("%d%d", &u, &v);
                    g[u].pb(v);
                    g[v].pb(u);
                }
                dfs1(1,0);
                dfs(1,0,1);
                for(int i=1; i<=n; i++)printf("%lld " , ans[i]);
                printf("\n");
                return 0;
}