CSU 1811 Tree Intersection（启发式合并）

思路：对于某一颗子树来说，只需要统计这颗子树拥有的颜色和这颗子树独有的颜色，两者一减就是答案，考虑一直DFS到子树，然后向上合并，合并的时候将结点数小的往大的合并，这样可以做到nlogn的复杂度

#include<bits/stdc++.h>
using namespace std;
const int maxn = 1e5+7;
struct Edge
{
	int v,id;
	Edge(){};
	Edge(int vv,int idd):v(vv),id(idd){};
};
vector<Edge>e[maxn];
int n,c[maxn],sum[maxn],res[maxn],ans[maxn];
map<int,int>cnt[maxn];    //cnt[u][i] 结点u颜色i有多少个
void dfs(int u,int fa,int id)
{
    cnt[u][c[u]]=1;
	ans[u]=cnt[u][c[u]]<sum[c[u]]?1:0;
	for(int i = 0;i<e[u].size();i++)
	{
		int v = e[u][i].v;
		int id = e[u][i].id;
		if(v==fa)continue;
		dfs(v,u,id);
        if(cnt[u].size()<cnt[v].size())
		{
			swap(ans[u],ans[v]);
			swap(cnt[u],cnt[v]);
		}
		for(map<int,int>::iterator it = cnt[v].begin();it!=cnt[v].end();it++)
		{
			int &num = cnt[u][it->first];
			if(num==0 && num+it->second<sum[it->first])
				++ans[u];
            else if (num && num+it->second==sum[it->first])
				--ans[u];
			num+=it->second;
		}
	}
	res[id]=ans[u];
}
int main()
{
	while(scanf("%d",&n)!=EOF)
	{
		memset(sum,0,sizeof(sum));
		for(int i = 0;i<=n;i++)
			e[i].clear(),cnt[i].clear();
		for(int i = 1;i<=n;i++)
			scanf("%d",&c[i]),sum[c[i]]++;
		for(int i = 1;i<=n-1;i++)
		{
			int u,v;
			scanf("%d%d",&u,&v);
			e[u].push_back(Edge(v,i));
			e[v].push_back(Edge(u,i));
		}
		dfs(1,-1,0);
        for(int i = 1;i<=n-1;i++)
			printf("%d\n",res[i]);
	} 
}

Description

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 _iand 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).

Output

For 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