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

600E - Lomsat gelral

题意

给出一颗以 1 为根的树，每个点有颜色，如果某个子树上某个颜色出现的次数最多，则认为它在这课子树有支配地位，一颗子树上，可能有多个有支配的地位的颜色，对每颗子树分别求有支配地位的颜色的和（把颜色这个权值相加）。

分析

树上启发式合并模板题。

参考blog1

参考blog2

复杂度证明

如果暴力去搜索，显然是 \(O(n^2)\) 的算法，可以考虑优化，当我们搜索到节点 u 时，最后去搜索 u 的子节点中子树节点数量最大的子节点（树链剖分求出重儿子），并保留这个子节点所在子树的状态（颜色数量信息），这样在更新贡献的时候可以直接跳过它了。复杂度 \(O(nlogn)\)。

code

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = 2e5 + 10;
int n;
int fa[MAXN], son[MAXN], dep[MAXN], siz[MAXN];
int col[MAXN];
int cnt, head[MAXN];
struct Edge {
    int to, next;
}e[MAXN];
void addedge(int u, int v) {
    e[cnt].to = v;
    e[cnt].next = head[u];
    head[u] = cnt++;
}
void dfs(int u) {
    siz[u] = 1; son[u] = 0;
    for(int i = head[u]; ~i; i = e[i].next) {
        if(e[i].to != fa[u]) {
            fa[e[i].to] = u;
            dep[e[i].to] = dep[u] + 1;
            dfs(e[i].to);
            if(siz[e[i].to] > siz[son[u]]) son[u] = e[i].to;
            siz[u] += siz[e[i].to];
        }
    }
}
int vis[MAXN];
ll sum, mx, C[MAXN];
ll ans[MAXN];
void update(int u, int c) {
    C[col[u]] += c;
    if(c > 0 && C[col[u]] >= mx) {
        if(C[col[u]] > mx) { sum = 0; mx = C[col[u]]; }
        sum += col[u];
    }
    for(int i = head[u]; ~i; i = e[i].next) {
        if(e[i].to != fa[u] && !vis[e[i].to]) update(e[i].to, c);
    }
}
void dfs1(int u, int flg) {
    for(int i = head[u]; ~i; i = e[i].next) {
        if(e[i].to != fa[u] && e[i].to != son[u]) dfs1(e[i].to, 1);
    }
    if(son[u]) {
        dfs1(son[u], 0);
        vis[son[u]] = 1;
    }
    update(u, 1); ans[u] = sum;
    if(son[u]) vis[son[u]] = 0;
    if(flg) {
        update(u, -1);
        sum = 0; mx = 0;
    }
}
int main() {
    memset(head, -1, sizeof head);
    scanf("%d", &n);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &col[i]);
    }
    for(int i = 1; i < n; i++) {
        int x, y;
        scanf("%d%d", &x, &y);
        addedge(x, y);
        addedge(y, x);
    }
    dfs(1);
    dfs1(1, -1);
    for(int i = 1; i <= n; i++) {
        printf("%I64d%c", ans[i], i == n ? '\n' : ' ');
    }
    return 0;
}