CSU 1811 Tree Intersection 平衡树启发式合并

题目链接：http://acm.csu.edu.cn/csuoj/problemset/problem?pid=1811

题意：每个点有一个颜色，删除一条边后，求这条边两边的点集的颜色交的个数。

解法：平衡树启发式合并。那么可以直接用map对每个子树的信息进行记录，然后回溯上去的时候父节点加上子节点信息

的时候，把小的往大的加然后每次新出来一个颜色，就+1，如果颜色满了，就−1每次新搜一个节点的时候，如果子节点

更大，那么ans也要赋值成子节点的ans 想一下就知道了，知道启发式合并之后，难点就是ans的变化了

#include <bits/stdc++.h>
using namespace std;
const int maxn=100010;
struct edge{
    int v,next,id;
}E[maxn*2];
int head[maxn],ans[maxn],sum[maxn],c[maxn],edgecnt;
map <int,int> cnt[maxn];
void init(){
    memset(head,-1,sizeof(head));
    edgecnt=0;
}
void add(int u,int v,int id){
    E[edgecnt].v=v,E[edgecnt].next=head[u],E[edgecnt].id=id,head[u]=edgecnt++;
}
void dfs(int u, int fa, int id)
{
    cnt[u][c[u]]=1;
    if(sum[c[u]]>1) ans[id]++;
    for(int i=head[u]; i+1; i=E[i].next){
        int v=E[i].v;
        if(v==fa) continue;
        dfs(v,u,E[i].id);
        if(cnt[u].size()<cnt[v].size()){
            //swap(cnt[u],cnt[v]);
            cnt[u].swap(cnt[v]);
            ans[id]=ans[E[i].id];
        }
        for(map<int,int>::iterator it=cnt[v].begin(); it!=cnt[v].end(); it++){
            int x=it->first,y=it->second;
            if(cnt[u].count(x)){
                cnt[u][x]+=y;
                if(cnt[u][x]==sum[x]) ans[id]--;
            }
            else{
                cnt[u][x]+=y;
                if(cnt[u][x]<sum[x]) ans[id]++;
            }
        }
    }
}
int main()
{
    int n;
    while(~scanf("%d", &n)){
        memset(ans, 0, sizeof(ans));
        memset(sum, 0, sizeof(sum));
        for(int i=0; i<maxn; i++) cnt[i].clear();
        for(int i=1; i<=n; i++) scanf("%d", &c[i]);
        for(int i=1; i<=n; i++) sum[c[i]]++;
        init();
        for(int i=1; i<n; i++){
            int u,v;
            scanf("%d%d",&u,&v);
            add(u,v,i);
            add(v,u,i);
        }
        dfs(1,-1,0);
        for(int i=1; i<n; i++) printf("%d\n", ans[i]);
    }
    return 0;
}