Tree Intersection(dsu)

Tree Intersection

题目：

枚举每一条边，把边拆开，分成两颗树，求两棵树相同颜色的种数。
树的大小$1e5$

题解：

两棵树用减法来分析
子树，整棵树$-$子树。
然后dsu就好了。

代码：

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+9;
int n,c[N];
struct Edge{
    int v,nxt,id;
}e[N<<1];
int head[N],cnt;
int all[N],ans[N];
inline void add(int u,int v,int id){
    e[cnt]=(Edge){v,head[u],id};
    head[u]=cnt++;
}
int siz[N],son[N];
void dfs(int u,int fa){
    siz[u]=1;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(e1.v==fa)continue;
        dfs(e1.v,u);
        siz[u]+=siz[e1.v];
        if(siz[son[u]]<siz[e1.v])son[u]=e1.v;
    }
}
int he[N],sum,nowson;
void calc(int u,int fa,int val){
 //   cout<<"calc\n";
    he[c[u]]+=val;                 //对于sum只考虑val=1的情况即可
    if(he[c[u]]==1)sum++;
    if(he[c[u]]==all[c[u]])sum--;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(e1.v==fa||e1.v==nowson)continue;
        calc(e1.v,u,val);
    }
}
void dsu(int u,int fa,int id,bool kep){
 //   cout<<"dus "<<u<<" "<<fa<<" "<<id<<" "<<kep<<endl;
    int sid;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(e1.v==son[u])sid=e1.id;
        if(e1.v==fa||e1.v==son[u])continue;
        dsu(e1.v,u,e1.id,0);
    }
    if(son[u])dsu(son[u],u,sid,1),nowson=son[u];
    calc(u,fa,1);nowson=0;
    ans[id]=sum;
    if(!kep){
        calc(u,fa,-1);
        sum=0;
    }
}
int main(){
  //  freopen("tt.in","r",stdin),freopen("tt.out","w",stdout);
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>n;
    memset(head,-1,sizeof(head));
    for(int i=1;i<=n;i++)cin>>c[i],all[c[i]]++;
    for(int i=1,u,v;i<n;i++){
        cin>>u>>v;
        add(u,v,i),add(v,u,i);
    }
    dfs(1,0);
    dsu(1,0,0,1);
    for(int i=1;i<n;i++)cout<<ans[i]<<endl;
    return 0;
}