启发式合并CSU

F - Tree Intersection CSU - 1811

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(b_i,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 ⁵).

The second line contains n integers c ₁,c ₂,…,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 ₁,R ₂,…,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

题意：给一棵带有颜色的树，去掉某一条边后，他就变成了两个连通块，求这两个连通块的颜色交集的大小。（对每条边求一次）
思路：首先我们可以预处理出来整棵树的颜色，去掉一条边，就相当于一颗子树独立出来了，就变成了一颗子树，和剩余的树，可以将子树的颜色处理出来，然后剩余的那棵树的颜色也求出来了（等于整树 减去 子树）。
子树的话我们就可以用启发性合并去求颜色。但是要注意一下细节，看下代码吧。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
using namespace std;
typedef long long LL;
const int maxn=1e5+5;
int n,m;
int mx,big,id,sum;
int col[maxn],deep[maxn],num[maxn],si[maxn],hson[maxn],cnt[maxn],ans[maxn];
int l[maxn],r[maxn];
vector<int>G[maxn];

void findhson(int x,int fa,int dep)//找到所有的重儿子，并求出其深度
{
    si[x]=1;
    deep[x]=dep;
    int len=G[x].size();
    for(int i=0;i<len;i++)
    {
        int t=G[x][i];
        if(t!=fa)
        {
            findhson(t,x,dep+1);
            si[x]+=si[t];
            deep[t]=deep[x]+1;
            if(si[t]>si[hson[x]])
                hson[x]=t;
        }
    }
}
void cal(int x,int fa,int val)
{
    cnt[col[x]]+=val;
    if(cnt[col[x]]==1)//如果新出现一种颜色就说明重合的数量多了一个
        sum++;
    if(cnt[col[x]]==num[col[x]])/*如果这个这个子树的该颜色数量等于整棵树的该
颜色数量，就说明两棵子树的重合数量减一*/
        sum--;
    int len=G[x].size();
    for(int i=0;i<len;i++)
    {
        int t=G[x][i];
        if(t!=fa && t!=big)
            cal(t,x,val);
    }
}
void dfs(int x,int fa,int flag)
{
    int len=G[x].size();
    for(int i=0;i<len;i++)
    {
        int t=G[x][i];
        if(t!=fa && t!=hson[x])
            dfs(t,x,0);
    }
    if(hson[x])
    {
        dfs(hson[x],x,1);
        big=hson[x];
    }
    cal(x,fa,1);
    big=0;
    ans[x]=sum;
    if(!flag)
    {
        cal(x,fa,-1);
        sum=0;
    }
}
int main()
{
    while(scanf("%d",&n)!=EOF)
    {
        big=0;sum=0;
        for(int i=1;i<=n;i++)
        {
            G[i].clear();
            num[i]=si[i]=hson[i]=cnt[i]=ans[i]=0;
        }
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&col[i]);
            num[col[i]]++;//记录每个颜色的数量
        }
        int x,y;
        for(int i=1;i<n;i++)
        {
            scanf("%d %d",&x,&y);
            G[x].push_back(y);
            G[y].push_back(x);
            l[i]=x;r[i]=y;//记录其子节点和父亲节点
        }
        findhson(1,0,1);
        dfs(1,0,1);
        for(int i=1;i<n;i++)
        {
            int j;
            if(deep[l[i]]>deep[r[i]])//这个看我等会传的图就知道为什么要比较一下了
                j=l[i];
            else
                j=r[i];
            printf("%d
",ans[j]);
        }
    }
    return 0;
}

为什么要比较深度？因为我们根固定，我们启发式合并只会求方框内的那个子树的情况，而并不会求剩余的那块的子树的情况。所以在计算的时候选择深度高的那颗子树去进行计算才是正确的。