Codeforces 600E dsu on tree 启发式合并

题目：传送门

题意：

给一颗树，计算所有结点的子树，颜色最多的结点的 颜色种类的sum和，子树的定义是当前结点和其所有的子节点。

题解：

我们可以想到一种暴力的解法，就是O(n)遍历所有的结点，然后dfs搜索它的所有子节点，查看颜色出现最多的结点，用一个ans数组维护答案，最后输出ans。显然，这样做的复杂度为O(n^2^)
我们这里用到dsu on tree算法，关于这个算法的详解，可以参考：http://codeforces.com/blog/entry/44351

AC 代码：

#include <bits/stdc++.h>
#define debug(x) cout<<#x<<" = "<<x<<endl;
#define INF 0x3f3f3f3f
using namespace std;

const int maxn = 100010;
vector<int> g[maxn];
int Max,col[maxn];
long long sum,ans[maxn];
/****************dsu on tree 算法************************/
int sz[maxn];
void getsz(int v, int p) {//计算每个结点有多少个子节点
    sz[v] = 1;  
    for (auto u : g[v])
        if (u != p) {
            getsz(u, v);
            sz[v] += sz[u];
        }
}
int cnt[maxn];
bool big[maxn];
void add(int v, int p, int x) {
    cnt[ col[v] ] += x;
    if (cnt[col[v]] > Max)
        sum = col[v], Max = cnt[col[v]];
    else if (cnt[col[v]] == Max)
        sum += col[v];
    for (auto u : g[v])
        if (u != p && !big[u])
            add(u, v, x);
}
void dfs(int v, int p, int keep) {//启发式搜索
    int mx = -1, bigChild = -1;
    for (auto u : g[v])
        if (u != p && sz[u] > mx)
            mx = sz[u], bigChild = u;
    for (auto u : g[v])
        if (u != p && u != bigChild)
            dfs(u, v, 0);  
    if (bigChild != -1)
        dfs(bigChild, v, 1), big[bigChild] = 1;  
    add(v, p, 1);
    if (bigChild != -1)
        big[bigChild] = 0;
    ans[v] = sum;
    if (keep == 0)
        add(v, p, -1), Max = sum = 0;
}
/****************dsu on tree 算法************************/
int main(void) {
    int n, t1, t2;
    cin >> n;
    for (int i = 1; i <=n; i++)
        cin >> col[i];
    for (int i = 1; i < n; i++)
        cin >> t1 >> t2, g[t1].push_back(t2), g[t2].push_back(t1);
    getsz(1, 0);
    dfs(1, 0, 1);
    for (int i = 1; i <= n; i++)
        cout << ans[i] << " ";
    return 0;
}