DSU on tree 树上启发式合并

自为风月马前卒dalao的博客

首先介绍一下大概流程：

• 首先处理所有轻链。
• 如果有重链，再处理重链。注意重链的值不删除。
• 这样只需要把轻链的贡献算一下加上就好了，不需要处理重链。
• 最后，如果是轻链，就要删除其对贡献的影响。

void dfs(int u, int pre, int opt){
    for(int i = head[u]; i != -1; i = nxt[i]){
        int v = to[i];
        if(v == pre) continue;
        if(v != son[u]) dfs(v, u, 0);
    }
    if(son[u]) dfs(son[u], u, 1), tson = son[u];
    Addval(), tson = 0;
    if(!opt) Delval();
}

再换一个角度，对于每个点，如果它是父节点的重儿子，返回前就不需要删除计算的贡献，这样相当于重儿子对父节点的贡献已经算了。如果它是父节点的轻儿子，那么返回前就需要删除轻儿子对父节点的影响。也就是从低向上保留重儿子对当前点u的贡献，再自上向下把轻儿子的贡献都加上。

复杂度： \(O(nlog^{n})\)

---

例题: Cf600E

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<map>
#include<queue>
#include<vector>
#include<string>
#include<fstream>
using namespace std;
#define rep(i, a, n) for(int i = a; i <= n; ++ i);
#define per(i, a, n) for(int i = n; i >= a; -- i);
typedef long long ll;
const int N = 2e5 + 105;
const int mod = 1e9 + 7;
const double Pi = acos(- 1.0);
const ll INF = 1e18;
const int G = 3, Gi = 332748118;
ll qpow(ll a, ll b) { ll res = 1; while(b){ if(b & 1) res = (res * a) % mod; a = (a * a) % mod; b >>= 1;} return res; }
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }
// bool cmp(int a, int b){return a > b;}
//

int n, m, A, B;
int head[N], cnt = 0;
int to[N << 1], nxt[N << 1];
int son[N], siz[N],dfn[N];
int col[N]; ll ans[N], sum;
int num[N], Max = 0, tson;

void add(int u, int v){
    to[cnt] = v, nxt[cnt] = head[u], head[u] = cnt ++;
    to[cnt] = u, nxt[cnt] = head[v], head[v] = cnt ++;
}

void dfs1(int u, int pre){
    siz[u] = 1;
    int maxx = -1;
    for(int i = head[u]; i != -1; i = nxt[i]){
        int v = to[i];
        if(v == pre) continue;
        dfs1(v, u);
        siz[u] += siz[v];
        if(siz[v] > maxx){
            maxx = siz[v];
            son[u] = v;
        }
    }
}

void Addval(int u, int pre, int val){
    num[col[u]] += val;
    if(num[col[u]] > Max) Max = num[col[u]], sum = col[u];
    else if(num[col[u]] == Max) sum += col[u];
    for(int i = head[u]; i != -1; i = nxt[i]){
        int v = to[i];
        if(v == pre || v == tson) continue;
        Addval(v, u, val);
    }
}

void dfs(int u, int pre, int opt){
    for(int i = head[u]; i != -1; i = nxt[i]){
        int v = to[i];
        if(v == pre) continue;
        if(v != son[u]) dfs(v, u, 0);
    }
    if(son[u]) dfs(son[u], u, 1), tson = son[u];
    Addval(u, pre, 1), tson = 0;
    ans[u] = sum;
    if(!opt){
        Addval(u, pre, -1);
        sum = 0, Max = 0;
    }
}

int main()
{
    scanf("%d",&n);
    cnt = 0;
    for(int i = 0; i <= n; ++ i) head[i] = -1;
    for(int i = 1; i <= n; ++ i) scanf("%d",&col[i]);
    for(int i = 1; i < n; ++ i){
        int x, y; scanf("%d%d",&x,&y);
        add(x, y);
    }
    dfs1(1, 0);
    dfs(1, 0, 0);
    for(int i = 1; i <= n; ++ i){
        printf("%lld",ans[i]);
        if(i == n) printf("\n");
        else printf(" ");
    }
    return 0;
}