【图论】树上启发式合并

reference
浅谈树链剖分
算法学习笔记(86): 树上启发式合并
算法学习笔记(59): 重链剖分

代码模板

https://codeforces.com/problemset/problem/600/E

int color[N], cnt[N], maxn;
ll sum, ans[N];
vector<int> edges[N];
int dep[N], sz[N], hson[N];

void dfs0(int u, int fa = -1) {
    // 树链剖分：统计子树大小，并找出重儿子
    int mx = 0, size = 1;
    for (auto v : edges[u]) {
        if (v != fa) {
            dfs0(v, u);
            size += sz[v];
            if (sz[v] > mx) {
                hson[u] = v, mx = sz[v];
            }
        }
    }
    sz[u] = size;
}

// 单个节点的信息更改
void add(int u) {
    int n = ++cnt[color[u]];
    if (n > maxn) {
        maxn = n, sum = color[u];
    }
    else if (n == maxn) {
        sum += color[u];
    }
}
void del(int u) {
    --cnt[color[u]];
}

// 递归地添加/删除子树信息
void addsubtree(int u, int fa = -1) {
    add(u);
    for (auto v : edges[u]) {
        if (v != fa) addsubtree(v, u);
    }
}
void delsubtree(int u, int fa = -1) {
    del(u);
    for (auto v : edges[u]) {
        if (v != fa) delsubtree(v, u);
    }
}

void dfs(int u, int fa = -1, bool keep = -1) {
    // keep标志信息是否保留
    // 先遍历轻儿子
    for (auto v : edges[u]) {
        if (v != fa && v != hson[u]) {
            dfs(v, u, 0);
        }
    }
    // 最后遍历重儿子
    if (hson[u]) { dfs(hson[u], u, 1); }
    add(u);
    for (auto v : edges[u]) {
        if (v != fa && v != hson[u]) {
            addsubtree(v, u);
        }
    }   // 重儿子的信息没有删除，因此不必再添加
    ans[u] = sum;
    if (!keep) {
        delsubtree(u, fa);
        // memset(cnt, 0, sizeof cnt);
        maxn = sum = 0;
    }
}