dsu on tree模板

处理树上不带修改的离线子树问题

关键词：静态、离线、子树

题目：CF600E

#include
using namespace std;
typedef long long ll;
const int N = 1e5 + 20, M = 2e5 + 20;
int n, c[N], cnt[N];
int flag, maxc;
ll ans[N], sum;

int e[M], ne[M], h[M], idx;
void add(int u, int v){
    e[idx] = v, ne[idx] = h[u], h[u] = idx++;
}

// 找重儿子
int siz[N], son[N];
void dfs(int u, int f){
    siz[u] = 1;
    for(int i = h[u]; ~i; i = ne[i]){
        int v = e[i];
        if(v == f) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(siz[v] > siz[son[u]]){
            son[u] = v;
        }
    }
}

// 统计答案
void count(int u, int f, int val){
    cnt[c[u]] += val;
    if(cnt[c[u]] > maxc){
        maxc = cnt[c[u]];
        sum = c[u];
    }
    else if(cnt[c[u]] == maxc)
        sum += c[u];
    for(int i = h[u]; ~i; i = ne[i]){
        int v = e[i];
        if(v == f || v == flag) continue;
        count(v, u, val);
    }
}

// dsu on tree 模板
void dfs_dsuontree(int u, int f, bool keep){
    // 计算轻儿子及其子树 统计答案并删除贡献
    for(int i = h[u]; ~i; i = ne[i]){
        int v = e[i];
        if(v == f || v == son[u]) continue;
        dfs_dsuontree(v, u, false);
    }
    // 计算重儿子及其子树 统计答案不删贡献
    if(son[u]){
        dfs_dsuontree(son[u], u, true);
        flag = son[u];
    }
    // 暴力统计u及其所有轻儿子的贡献 合并到刚算出的重儿子信息里
    count(u, f, 1);
    flag = 0;
    ans[u] = sum;
    // 把需要删除贡献删除
    if(!keep){
        count(u, f, -1);
        sum = maxc = 0;
    }
}

int main(){
    memset(h, -1, sizeof h);
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) scanf("%d", &c[i]);
    for(int i = 1; i < n; i++){
        int u, v;
        scanf("%d%d", &u, &v);
        add(u, v);
        add(v, u);
    }
    dfs(1, -1);
    dfs_dsuontree(1, 0, false);
    for(int i = 1; i <= n; i++){
        printf("%lld ", ans[i]);
    }
    return 0;
}