CSU 1811 Tree Intersection (dsu on tree)

Tree Intersection

Bobo has a tree with n vertices numbered by 1,2,…,n and (n-1) edges. The i-th vertex has color c i, and the i-th edge connects vertices a i and b i.
Let C(x,y) denotes the set of colors in subtree rooted at vertex x deleting edge (x,y).
Bobo would like to know R_i which is the size of intersection of C(a i,b i) and C(bi,a i) for all 1≤i≤(n-1). (i.e. |C(a i,b i)∩C(b i,a i)|)

Input

The input contains at most 15 sets. For each set:
The first line contains an integer n (2≤n≤10 5).
The second line contains n integers c 1,c 2,…,c n (1≤c_i≤n).
The i-th of the last (n-1) lines contains 2 integers a i,b i (1≤a i,b i≤n).
OutputFor each set, (n-1) integers R 1,R 2,…,R n-1.

Sample Input

4
1 2 2 1
1 2
2 3
3 4
5
1 1 2 1 2
1 3
2 3
3 5
4 5

Sample Output

1
2
1
1
1
2
1

题意:

给一棵n个节点的树，每个节点有颜色，问对于每一条边，删去这条边之后形成的两个连通块的颜色种类交集有多少种颜色。

思路：

我想到的是把问题转化成：总颜色数-一端独有的颜色数-另一端独有的颜色数，
然后没写出来。

看别人把问题转化成了：子树颜色数-子树独有的颜色数，
因为只分为两个连通块，仔细想想这样确实是对的，这个思路妙啊！

而且这样就转化为子树问题，可以用dsu on tree解决

code：

#include<bits/stdc++.h>
using namespace std;
const int maxm=1e5+5;
int head[maxm],nt[maxm<<1],to[maxm<<1],w[maxm<<1],tot;//存边，其中w[i]是边的id
int sz[maxm],son[maxm],sid[maxm];//sid记录连接重儿子的边的id
int mark[maxm];//用于标记重儿子
int sum[maxm];//每种颜色的总个数
int cnt[maxm];//当前子树的颜色个数
int ans[maxm];//存答案
int a[maxm];//存颜色
int all,num;//all是当前子树的颜色种类，num是当前子树独有的颜色种类
int n;
void init(){
    for(int i=1;i<=n;i++){
        head[i]=0;
        sum[i]=0;
        cnt[i]=0;
    }
    tot=0;
    num=0;
    all=0;
}
void add(int x,int y,int z){
    tot++;nt[tot]=head[x];head[x]=tot;to[tot]=y;w[tot]=z;
}
void dfs(int x,int fa){//求重儿子
    sz[x]=1;
    son[x]=0;
    for(int i=head[x];i;i=nt[i]){
        int v=to[i];
        if(v==fa)continue;
        dfs(v,x);
        sz[x]+=sz[v];
        if(sz[v]>sz[son[x]]){
            son[x]=v;
            sid[x]=w[i];
        }
    }
}
void cal(int x,int fa,int change){
    cnt[a[x]]+=change;
    if(change>0){//add
        if(cnt[a[x]]==1)all++;
        if(cnt[a[x]]==sum[a[x]])num++;
    }else{//del
        if(cnt[a[x]]==0)all--;
        if(cnt[a[x]]==sum[a[x]]-1)num--;
    }
    for(int i=head[x];i;i=nt[i]){
        int v=to[i];
        if(v==fa||mark[v])continue;
        cal(v,x,change);
    }
}
void solve(int x,int fa,int kep,int id){
    for(int i=head[x];i;i=nt[i]){//先解决轻儿子并擦除
        int v=to[i];
        if(v==fa||v==son[x])continue;
        solve(v,x,0,w[i]);
    }
    if(son[x]){//解决重儿子并不擦除
        solve(son[x],x,1,sid[x]);
        mark[son[x]]=1;
    }
    cal(x,fa,1);//计算轻儿子贡献
    ans[id]=all-num;
    if(son[x]){
        mark[son[x]]=0;
    }
    if(!kep){
        cal(x,fa,-1);
        all=num=0;
    }
}
signed main(){
    while(scanf("%d",&n)!=EOF){
        init();
        for(int i=1;i<=n;i++){//输入颜色
            scanf("%d",&a[i]);
            sum[a[i]]++;//统计每种颜色总数
        }
        for(int i=1;i<n;i++){
            int a,b;
            scanf("%d%d",&a,&b);
            add(a,b,i);
            add(b,a,i);
        }
        dfs(1,-1);
        solve(1,-1,1,0);
        for(int i=1;i<n;i++){
            printf("%d\n",ans[i]);
        }
    }
    return 0;
}
//https://ac.nowcoder.com/acm/contest/1112/I