【题解】Lomsat gelral 树上启发式合并 CF – 600E

题目链接

树上启发式合并（dsu on tree）模板题

就是先轻重链剖分，dfs 算答案的时候，先去计算轻儿子的答案，为了防止儿子之间互相影响，每计算完一个儿子就要清空统计信息

然后 dfs 算重儿子答案，此时不清空统计信息，然后再去把所有轻儿子的统计信息算进来就能得到当前点的答案了

这样子每个点被统计的次数为其返祖链上的轻边数目，复杂度为 $O(n \log _ 2 n)$

代码：

#include <bits/stdc++.h>

#define NS (100005)

using namespace std;

typedef long long LL;

template<typename _Tp> inline void IN(_Tp &dig)
{
    char c; bool flag = 0; dig = 0;
    while (c = getchar(), !isdigit(c)) if (c == '-') flag = 1;
    while (isdigit(c)) dig = dig * 10 + c - '0', c = getchar();
    if (flag) dig = -dig;
}

int n, v[NS], sz[NS], hs[NS];

struct graph
{
    int head[NS], nxt[NS << 1], to[NS << 1], sz;
    graph() { memset(head, -1, sizeof(head)), sz = 0; }
    void push(int a, int b)
    {
        nxt[sz] = head[a], to[sz] = b, head[a] = sz++;
    }
    int operator [] (const int a) const { return to[a]; }
} g;

void dfs(int a, int fa)
{
    sz[a] = 1;
    for (int i = g.head[a]; ~i; i = g.nxt[i])
        if (g[i] != fa)
        {
            dfs(g[i], a), sz[a] += sz[g[i]];
            if (sz[g[i]] > sz[hs[a]]) hs[a] = g[i];
        }
}

int h, cnt[NS], mx;

LL sum, ans[NS];

void push(int a, int fa, int x)
{
    cnt[v[a]] += x;
    if (cnt[v[a]] > mx) mx = cnt[v[a]], sum = v[a];
    else if (cnt[v[a]] == mx) sum += v[a];
    for (int i = g.head[a]; ~i; i = g.nxt[i])
        if (g[i] != fa && g[i] != h)
            push(g[i], a, x);
}

void dsu(int a, int fa, bool t)
{
    for (int i = g.head[a]; ~i; i = g.nxt[i])
        if (g[i] != fa && g[i] != hs[a])
            dsu(g[i], a, 0);
    if (hs[a]) dsu(hs[a], a, 1), h = hs[a];
    push(a, fa, 1), ans[a] = sum, h = 0;
    if (!t) push(a, fa, -1), mx = 0;
}

int main(int argc, char const* argv[])
{
    IN(n);
    for (int i = 1; i <= n; i += 1) IN(v[i]);
    for (int i = 1, a, b; i < n; i += 1)
        IN(a), IN(b), g.push(a, b), g.push(b, a);
    dfs(1, 0), dsu(1, 0, 1), printf("%lld", ans[1]);
    for (int i = 2; i <= n; i += 1) printf(" %lld", ans[i]);
    putchar(10);
    return 0;
}