Codeforces 600E dsu on tree

题意

传送门 Codeforces 600E Lomsat gelral

题解

求各子树上数量最多的颜色之和。$dsuontree$ 求解，总时间复杂度 $O(N\log N)$。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 100005, maxe = 2 * maxn;
int N, mx, C[maxn], cnt[maxn];
ll sum, res[maxn];
int tot, head[maxn], to[maxe], nxt[maxe];
int sz[maxn], hs[maxn];
bool hv[maxn];

inline void add(int x, int y) { to[++tot] = y, nxt[tot] = head[x], head[x] = tot; }

void get_sz(int x, int f)
{
    sz[x] = 1;
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f)
            get_sz(y, x), sz[x] += sz[y];
    }
}

void add(int x, int f, int d)
{
    cnt[C[x]] += d;
    if (d > 0)
    {
        if (cnt[C[x]] > mx)
            mx = cnt[C[x]], sum = C[x];
        else if (cnt[C[x]] == mx)
            sum += C[x];
    }
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && !hv[y])
            add(y, x, d);
    }
}

void dfs(int x, int f, int keep)
{
    int hs = 0;
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && sz[y] > sz[hs])
            hs = y;
    }
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != hs && y != f)
            dfs(y, x, 0);
    }
    if (hs)
        dfs(hs, x, 1), hv[hs] = 1;
    add(x, f, 1);
    if (hs)
        hv[hs] = 0;
    res[x] = sum;
    if (!keep)
        add(x, f, -1), mx = sum = 0;
}

int main()
{
    scanf("%d", &N);
    for (int i = 1; i <= N; ++i)
        scanf("%d", C + i);
    for (int i = 1, x, y; i < N; ++i)
        scanf("%d%d", &x, &y), add(x, y), add(y, x);
    get_sz(1, 0);
    dfs(1, 0, 1);
    for (int i = 1; i <= N; ++i)
        printf("%lld%c", res[i], i == N ? '\n' : ' ');
    return 0;
}