CSUOJ1811 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(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

Hint

 题解：题意就是，给以一颗树n个节点，每个节点有一种颜色，然年后对于n-1条边，如果把一条边截断，让你求两颗子树有多少种相同的颜色，依次输入每一条边的答案。

启发式搜索，分别记录点和边的答案；如果点u和其子树某种颜色的数量已经等于总量了，那么对于该子树外的一部分，就没有该中颜色了，答案-1；如果小于总量，答案+1；

然后更新u节点该颜色的数量即可；

参考代码：

#include<bits/stdc++.h>
using namespace std;
#define clr(a,val) memset(a,val,sizeof (a))
#define pb push_back
#define fi first
#define se second
typedef long long ll;
const int maxn=1e5+10;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-') f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=(x<<3)+(x<<1)+ch-'0';ch=getchar();}
    return x*f;
}
struct Edge{
    int to,index,nxt;
} edge[maxn<<1];
int n,head[maxn<<1],tot;
int col[maxn],sum[maxn],ans[maxn],res[maxn<<1];//ans[u]表示u点及子节点的答案, res[edge]表示边的答案
map<int,int> cnt[maxn];//cnt[u][color] 表示u点子树color颜色有多少个节点

inline void Init()
{
    clr(head,-1);clr(sum,0); tot=0;
    for(int i=0;i<=n;++i) cnt[i].clear();
}

inline void addedge(int u,int v,int id)
{
    edge[tot].to=v;
    edge[tot].index=id;
    edge[tot].nxt=head[u];
    head[u]=tot++;
}

inline void dfs(int u,int fa,int id)
{
    cnt[u][col[u]]=1;
    ans[u] = cnt[u][col[u]]<sum[col[u]]?1:0;
    for(int e=head[u];~e;e=edge[e].nxt)
    {
        int v=edge[e].to;
        if(v==fa) continue;
        dfs(v,u,edge[e].index);
        if(cnt[u].size()<cnt[v].size())
        {
            swap(cnt[u],cnt[v]);
            swap(ans[u],ans[v]);
        }
        map<int,int>::iterator it;
        for(it=cnt[v].begin();it!=cnt[v].end();it++)
        {
            if(!cnt[u][(*it).fi] && (*it).se<sum[(*it).fi]) ++ans[u];
            else if(cnt[u][(*it).fi] && cnt[u][(*it).fi]+(*it).se==sum[(*it).fi]) --ans[u];
            cnt[u][(*it).fi]+=(*it).se;//加上子树的数量 
        }
    }
    res[id]=ans[u];
}

int main()
{
    while(~scanf("%d",&n))
    {
        Init();
        for(int i=1;i<=n;++i) col[i]=read(),sum[col[i]]++;
        for(int i=1;i<n;++i) 
        {
            int u=read(),v=read();
            addedge(u,v,i);addedge(v,u,i);
        }
        dfs(1,0,0);
        for(int i=1;i<n;++i) printf("%d
",res[i]);    
    }
    
    return 0;
}