树上启发式合并【学习博客例题】

学习博客

例题：

牛客练习赛：DongDong数颜色

代码：

#include<bits/stdc++.h>
#define lson sign << 1, l, mid
#define rson sign << 1 | 1, mid + 1, r
#define lowbit(x) (x&(-x))
typedef long long ll;
using namespace std;
const int INF = 0x3f3f3f3f;
const int maxn = 1e5 + 5;

int n, m, color[maxn];
int sizes[maxn], son[maxn], now;
int cnt[maxn], Ans;
int ans[maxn];
vector<int>load[maxn];

void dfs1(int s, int pre) {
    sizes[s] = 1;
    for(int &e : load[s]) {
        if(e == pre)
            continue;
        dfs1(e, s);
        sizes[s] += sizes[e];
        if(sizes[e] > sizes[son[s]])
            son[s] = e;
    }
}
void add(int s, int pre, int flag) {
    cnt[color[s]] += flag;
    if(cnt[color[s]] == 1 && flag == 1)
        Ans++;
    if(cnt[color[s]] == 0 && flag == -1)
        Ans--;
    for(int &e : load[s]) {
        if(e == pre || e == now)
            continue;
        add(e, s, flag);
    }
}
void dfs2(int s, int pre, int flag) {
    for(int &e : load[s]) {
        if(e == pre || e == son[s])
            continue;
        dfs2(e, s, 0);
        ///记录轻儿子的贡献
        ///消除其影响
    }
    ///记录重儿子的贡献,保留其影响
    if(son[s])
        dfs2(son[s], s, 1);
    ///记录轻儿子和自身的贡献
    now = son[s];   ///不记录重儿子
    add(s, pre, 1);
    ///更新答案
    ans[s] = Ans;
    if(!flag) {
        now = 0;    ///消去整棵数贡献
        add(s, pre, -1);
    }
}
int main() {
    int s, e;
    scanf("%d %d", &n, &m);
    for(int i = 1; i <= n; i++)
        scanf("%d", &color[i]);
    for(int i = 1; i < n; i++) {
        scanf("%d %d", &s, &e);
        load[s].push_back(e);
        load[e].push_back(s);
    }
    dfs1(1, 0);
    dfs2(1, 0, 0);
    while(m--) {
        scanf("%d", &s);
        printf("%d\n", ans[s]);
    }
    return 0;
}

CF：Lomsat gelral

题意：给出一个树，求出每个节点的子树中出现次数最多的颜色的颜色和。

代码：

#include<bits/stdc++.h>
#define lson sign << 1, l, mid
#define rson sign << 1 | 1, mid + 1, r
#define lowbit(x) (x&(-x))
typedef long long ll;
using namespace std;
const int INF = 0x3f3f3f3f;
const int maxn = 1e5 + 5;

int n, color[maxn];
int sizes[maxn], son[maxn];
vector<int>load[maxn];

void dfs1(int s, int pre) {
    sizes[s] = 1;
    for(int &e : load[s]) {
        if(e == pre)
            continue;
        dfs1(e, s);
        sizes[s] += sizes[e];
        if(sizes[e] > sizes[son[s]])
            son[s] = e;
    }
}
ll cnt[maxn], anscnt, sum, now;
ll ans[maxn];
void add(int s, int pre, int val) {
    cnt[color[s]] += val;
    if(cnt[color[s]] > anscnt) {
        anscnt = cnt[color[s]];
        sum = color[s];
    } else if(cnt[color[s]] == anscnt)
        sum += color[s];
    for(int &e : load[s]) {
        if(e == pre || e == now)
            continue;
        add(e, s, val);
    }
}
void dfs2(int s, int pre, int flag) {
    for(int &e : load[s]) {
        if(e == pre || e == son[s])
            continue;
        dfs2(e, s, 0);
    }
    if(son[s])
        dfs2(son[s], s, 1);
    now = son[s];
    add(s, pre, 1);
    ans[s] = sum;
    if(!flag) {
        now = 0;
        add(s, pre, -1);
        anscnt = sum = 0;
    }
}
int main() {
    cin >> n;
    int s, e;
    for(int i = 1; i <= n; i++)
        cin >> color[i];
    for(int i = 1; i < n; i++) {
        cin >> s >> e;
        load[s].push_back(e);
        load[e].push_back(s);
    }
    dfs1(1, 0);
    dfs2(1, 0, 0);
    for(int i = 1; i <= n; i++)
        cout<<ans[i]<<' ';
    return  0;
}

牛客：Tree Intersection

题意：给你一棵树，每个点有一种颜色，问去掉每条边后，剩下两个连通块颜色种类的交集数量

答案：子树颜色种类-子树独有颜色种类

代码:

#include<bits/stdc++.h>
#define lson sign << 1, l, mid
#define rson sign << 1 | 1, mid + 1, r
#define lowbit(x) (x&(-x))
typedef long long ll;
using namespace std;
const int INF = 0x3f3f3f3f;
const int maxn = 1e5 + 5;

int n, c[maxn], sum[maxn], fa[maxn];
int sz[maxn], son[maxn];
vector<int>load[maxn];

void dfs1(int s, int pre) {
    sz[s] = 1;
    fa[s] = pre;
    for(auto e : load[s]) {
        if(e == pre)
            continue;
        dfs1(e, s);
        sz[s] += sz[e];
        if(sz[e] > sz[son[s]])
            son[s] = e;
    }
}
int ans[maxn], cnt[maxn], no, all, only;

void add(int s, int pre, int val) {
    cnt[c[s]] += val;
    if(val == -1) {
        if(cnt[c[s]] == 0)
            all--;
        if(cnt[c[s]] == sum[c[s]] - 1)
            only--;
    } else {
        if(cnt[c[s]] == 1)
            all++;
        if(cnt[c[s]] == sum[c[s]])
            only++;
    }
    for(auto e : load[s]) {
        if(e == pre || e == no)
            continue;
        add(e, s, val);
    }
}
void dfs2(int s, int pre, int flag) {
    for(auto e : load[s]) {
        if(e == pre || e == son[s])
            continue;
        dfs2(e, s, 0);
    }
    if(son[s])
        dfs2(son[s], s, 1);
    no = son[s];
    add(s, pre, 1);
    ans[s] = all - only;
    if(!flag) {
        no = 0;
        add(s, pre, -1);
    }
}
int s[maxn], e[maxn];
void init() {
    for(int i = 1; i <= n; i++) {
        load[i].clear();
        cnt[i] = sum[i] = son[i] = 0;
    }
}
int main() {
    while(cin >> n) {
        init();
        for(int i = 1; i <= n; i++) {
            cin >> c[i];
            sum[c[i]]++;
        }
        for(int i = 1; i < n; i++) {
            cin >> s[i] >> e[i];
            load[s[i]].push_back(e[i]);
            load[e[i]].push_back(s[i]);
        }
        dfs1(1, 0);
        dfs2(1, 0, 0);
        for(int i = 1; i < n; i++) {
            if(fa[s[i]] == e[i])
                cout << ans[s[i]] << endl;
            else
                cout << ans[e[i]] << endl;
        }
    }
}