cf600E. Lomsat gelral（树上启发式合并）

Posted 2020-11-23 heyuhhh

传送门

题意：
求子树众数。

思路：
(dsu on tree)模板题，用一个桶记录即可。
感觉(dsu on tree)这个算法的涉及真是巧妙呀，保留重链的信息，不断暴力轻边，并且不断在子树内递归下去。又由于轻边数量不会超过(O(logn))，所以总的时间复杂度控制在(O(nlogn))。

/*
 * Author:  heyuhhh
 * Created Time:  2019/11/13 15:09:25
 */
#include <bits/stdc++.h>
#define MP make_pair
#define fi first
#define se second
#define sz(x) (int)(x).size()
#define all(x) (x).begin(), (x).end()
#define INF 0x3f3f3f3f
#define Local
#ifdef Local
  #define dbg(args...) do { cout << #args << " -> "; err(args); } while (0)
  void err() { std::cout << '
'; }
  template<typename T, typename...Args>
  void err(T a, Args...args) { std::cout << a << ' '; err(args...); }
#else
  #define dbg(...)
#endif
void pt() {std::cout << '
'; }
template<typename T, typename...Args>
void pt(T a, Args...args) {std::cout << a << ' '; pt(args...); }
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
//head
const int N = 1e5 + 5;

int n;
int c[N];
struct Edge{
    int v, next;
}e[N << 1];
int head[N], tot;
void adde(int u, int v) {
    e[tot].v = v; e[tot].next = head[u]; head[u] = tot++;
}
int Son[N], sz[N];
void dfs(int u, int fa) {
    int mx = 0, son = 0;
    sz[u] = 1;
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v;
        if(v == fa) continue;
        dfs(v, u);
        sz[u] += sz[v];
        if(sz[v] > mx) {
            mx = sz[v], son = v;
        }
    }
    Son[u] = son;
}
ll sum, ans[N];
int son, Max;
int cnt[N];
void add(int u, int fa, int val) {
    cnt[c[u]] += val;
    if(cnt[c[u]] > Max) Max = cnt[c[u]], sum = c[u];
    else if(cnt[c[u]] == Max) sum += c[u];
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v;
        if(v == fa || v == son) continue;
        add(v, u, val);
    }
}
void dfs2(int u, int fa, int op) {
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v;
        if(v == fa || v == Son[u]) continue;
        dfs2(v, u, 0);
    }
    if(Son[u]) dfs2(Son[u], u, 1);
    son = Son[u];
    add(u, fa, 1);
    ans[u] = sum;
    son = 0;
    if(!op) add(u, fa, -1), Max = sum = 0;
}
void run(){
    memset(head, -1, sizeof(head));
    for(int i = 1; i <= n; i++) cin >> c[i];
    for(int i = 1; i < n; i++) {
        int u, v; cin >> u >> v;
        adde(u, v); adde(v, u);
    }
    dfs(1, 0);
    dfs2(1, 0, 1);
    for(int i = 1; i <= n; i++) cout << ans[i] << " 
"[i == n];
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    cout << fixed << setprecision(20);
    while(cin >> n) run();
    return 0;
}