Codeforces 600E Lomsat gelral （启发式合并）

Codeforces 600E Lomsat gelral （启发式合并）

题目大意：
一棵树，每个点有一个颜色，每一次询问在u节点为根的子树中，颜色出现次数最多的那些颜色的和。
考虑启发式合并，每个节点开map，cnt[u][i]表示u节点为根的子树上，i颜色在出现多少次，sum[u][i]表示在当前子树中，恰好出现了i次的颜色和。
每次把小的插入到大的里边，修改cnt，sum
合并复杂度O(nlogn), map的复杂度O(logn), 总复杂度O(nlog2n)

#include <map>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <vector>
#include <algorithm>
#define LL long long
#define N 100010
using namespace std;

struct Edge{
    int to, nxt;
}ed[N << 1];

int n, idc=0;
int head[N];
int col[N];
LL ans[N];
map<int, int> cnt[N];
map<int, LL> sum[N];

void adde(int u,int v){
    ed[++idc].to = v;
    ed[idc].nxt = head[u];
    head[u] = idc;
}

void dfs(int u,int fa){
    cnt[u].clear(); sum[u].clear();
    cnt[u][col[u]] = 1;
    sum[u][1] = col[u];
    for(int i=head[u]; i; i=ed[i].nxt){
        int v = ed[i].to;
        if(v == fa) continue;
        dfs(v, u);
        if(cnt[u].size() < cnt[v].size()){
            swap(cnt[u], cnt[v]);
            swap(sum[u], sum[v]);
        }
        for(map<int,int>::iterator it=cnt[v].begin(); it!=cnt[v].end(); it++){
            int x = it->first, y = it->second;
            if( cnt[u].count(x) ){
                sum[u][cnt[u][x]] -= x;
                cnt[u][x] += y;
                sum[u][cnt[u][x]] += x;
            }
            else{
                cnt[u][x] = y;
                sum[u][cnt[u][x]] += x;
            }
        }
    }
    map<int, LL>::iterator it = sum[u].end(); it--;
    ans[u] = it->second;//最大次数颜色和 
}

void init(){
    idc = 0;
    memset(head, 0, sizeof(head));
    memset(ed, 0, sizeof(ed));
}

int main(){
    while( scanf("%d", &n) != EOF ){
        init();
        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);
            adde(u, v); adde(v, u);
        }
        dfs(1, 1);
        for(int i=1; i<=n; i++) printf("%I64d ", ans[i]);
    }
    return 0;
}