典题DSU on tree求助!悬关一个
查看原帖
典题DSU on tree求助!悬关一个
519384
Link_Cut_Y楼主2023/5/27 15:15

map 写的启发式合并。样例已过。

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <map>

using namespace std;

const int N = 100010, M = N << 1;

int h[N], e[M], ne[M], idx, c[N], n;
int ans[N], max_val[N], sz[N], son[N];
map<int, int> s[N];

void add(int a, int b) {
    e[ ++ idx] = b, ne[idx] = h[a], h[a] = idx;
}
void init(int u, int father) { // 预处理重儿子
    sz[u] = 1;
    for (int i = h[u]; i; i = ne[i]) {
        int v = e[i]; if (v == father) continue;
        init(v, u); sz[u] += sz[v];
        if (sz[v] > sz[son[u]]) son[u] = v;
    }
}
void dfs(int u, int father) {
    if (son[u]) {
        dfs(son[u], u);
        s[u] = s[son[u]];
        max_val[u] = max_val[son[u]];
        ans[u] = ans[son[u]];
        s[son[u]].clear();
    }
    for (int i = h[u]; i; i = ne[i]) {
        int v = e[i];
        if (v == father or v == son[u]) continue;
        dfs(v, u);
        for (auto [x, y] : s[v]) {
            s[u][x] += y;
            if (s[u][x] == max_val[u])
                ans[u] += x;
            else if (s[u][x] > max_val[u])
                ans[u] = x, max_val[u] = s[u][x];
        }
        s[v].clear();
    }
    s[u][c[u]] ++ ;
    if (s[u][c[u]] == max_val[u])
        ans[u] += c[u];
    else if (s[u][c[u]] > max_val[u])
        ans[u] = c[u]; max_val[u] = s[u][c[u]];
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i ++ )
        scanf("%d", &c[i]);
    for (int i = 1; i < n; i ++ ) {
        int a, b; scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }
    init(1, -1); dfs(1, -1);
    for (auto i = 1; i <= n; i ++ )
        printf("%d ", ans[i]);
    return 0;
}
2023/5/27 15:15
加载中...