CF 600 E. Lomsat gelral

E. Lomsat gelral

http://codeforces.com/contest/600/problem/E

题意：

　　求每个子树内出现次数最多的颜色（如果最多的颜色出现次数相同，将颜色编号求和）。

分析：

　　dsu on tree。

　　这个可以解决一系列不带修改的子树查询问题。

　　考虑暴力的思路：就是枚举每个子树，计算每个颜色出现的个数。统计答案。

　　dsu on tree：最后一个子树枚举计算完了，它的贡献可以保留，然后用其它的子树去合并。（最后一棵子树是最大的）。现在的复杂度就是nlogn了。

代码：

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<iostream>
#include<cctype>
#include<set>
#include<vector>
#include<queue>
#include<map>
using namespace std;
typedef long long LL;

inline int read() {
    int x=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
    for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
}

const int N = 100005;

int head[N], nxt[N << 1], to[N << 1], En;
int fa[N], siz[N], son[N], cnt[N], col[N], Mx;
LL ans[N], Sum;

void add_edge(int u,int v) {
    ++En; to[En] = v; nxt[En] = head[u]; head[u] = En;
    ++En; to[En] = u; nxt[En] = head[v]; head[v] = En;
}

void dfs(int u) {
    siz[u] = 1;
    for (int i=head[u]; i; i=nxt[i]) {
        int v = to[i];
        if (v == fa[u]) continue;
        fa[v] = u;
        dfs(v);
        siz[u] += siz[v];
        if (!son[u] || siz[son[u]] < siz[v]) son[u] = v;
    }
}

void add(int u) {
    cnt[col[u]] ++;
    if (cnt[col[u]] > Mx) Mx = cnt[col[u]], Sum = col[u];
    else if (cnt[col[u]] == Mx) Sum += col[u];
}
void Calc(int u) {
    add(u);
    for (int i=head[u]; i; i=nxt[i]) 
        if (to[i] != fa[u]) Calc(to[i]);
}
void Clear(int u) {
    cnt[col[u]] --;
    for (int i=head[u]; i; i=nxt[i]) 
        if (to[i] != fa[u]) Clear(to[i]);
}

void solve(int u,bool c) {
    for (int i=head[u]; i; i=nxt[i]) 
        if (to[i] != fa[u] && to[i] != son[u]) solve(to[i], 0);
    if (son[u]) solve(son[u], 1);
    add(u);
    for (int i=head[u]; i; i=nxt[i]) 
        if (to[i] != fa[u] && to[i] != son[u]) Calc(to[i]);
    ans[u] = Sum;
    if (!c) Clear(u), Mx = 0, Sum = 0;
}

int main() {
    int n = read();
    for (int i=1; i<=n; ++i) col[i] = read();
    for (int i=1; i<n; ++i) {
        int u = read(), v = read(); 
        add_edge(u, v);
    }
    dfs(1);
    solve(1, 1);
    for (int i=1; i<=n; ++i) printf("%I64d ",ans[i]);
    return 0;
}