HDU6504 Problem E. Split The Tree【dsu on tree】

Problem E. Split The Tree

Problem Description
You are given a tree with n vertices, numbered from 1 to n. ith vertex has a value wi
We define the weight of a tree as the number of different vertex value in the tree.
If we delete one edge in the tree, the tree will split into two trees. The score is the sum of these two trees’ weights.
We want the know the maximal score we can get if we delete the edge optimally

给出一棵根为(1)的树，每个节点有权值，现在断开一条边，把树一分为二，记一棵树的贡献为这棵树里所有节点中不同权值的数量，现在要计算断开边之后两棵树的贡献的和的最大值

现在有一棵子树，为了得到答案，我们只要知道它里面出现过的各个权值有多少个，记原树中各个权值出现的数量是(cnt[i])，分开之后的一棵子树的各个权值出现的数量使(dcnt[i])，只要出现(dcnt[i] e cnt[i] and dcnt[i] e 0)，对答案的贡献就(+1)，这个可以(O(1))处理，那么就可以树上启发式合并了

//#pragma GCC optimize("O3")
//#pragma comment(linker, "/STACK:1024000000,1024000000")
#include<bits/stdc++.h>
using namespace std;
function<void(void)> ____ = [](){ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);};
const int MAXN = 1e5+7;
int n,w[MAXN],ret,tmp,cnt[MAXN],dcnt[MAXN],son[MAXN],sz[MAXN],st[MAXN],ed[MAXN],id[MAXN],dfn;
vector<int> G[MAXN];

void dfs(int u){
    sz[u] = 1; son[u] = 0;
    st[u] = ++dfn; id[dfn] = u;
    for(int v : G[u]){
        dfs(v);
        sz[u] += sz[v];
        if(sz[v]>sz[son[u]]) son[u] = v;
    }
    ed[u] = dfn;
}
void update(int val, int inc){
    if(inc==1){
        if(!dcnt[val]) tmp++;
        dcnt[val]++;
        if(dcnt[val]==cnt[val]) tmp--;
    }
    else{
        if(dcnt[val]==cnt[val]) tmp++;
        dcnt[val]--;
        if(!dcnt[val]) tmp--;
    }
}
void search(int u, bool clear){
    for(int v : G[u]) if(v!=son[u]) search(v,true);
    if(son[u]) search(son[u],false);
    for(int v : G[u]) if(v!=son[u]) for(int i = st[v]; i <= ed[v]; i++) update(w[id[i]],1);
    update(w[u],1);
    ret = max(ret,tmp);
    if(clear) for(int i = st[u]; i <= ed[u]; i++) update(w[id[i]],-1);
}
void solve(){
    dfn = 0;
    for(int i = 1; i <= n; i++) G[i].clear();
    for(int i = 2; i <= n; i++){
        int par; scanf("%d",&par);
        G[par].push_back(i);
    }
    tmp = 0;
    for(int i = 1; i <= n; i++){
        scanf("%d",&w[i]);
        if(!cnt[w[i]]) tmp++;
        cnt[w[i]]++;
    }
    ret = tmp;
    dfs(1);
    search(1,true);
    printf("%d
",ret);
    for(int i = 1; i <= n; i++) cnt[w[i]]--;
}
int main(){
    while(scanf("%d",&n)!=EOF) solve();
    return 0;
}