F - Tree Intersection

最新推荐文章于 2024-07-04 19:37:51 发布

evilwind2000

最新推荐文章于 2024-07-04 19:37:51 发布

阅读量99

点赞数

本文链接：https://blog.csdn.net/evilwind2000/article/details/103938130

版权

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 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).
Output
For 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
Hint
用map记录颜色，额外开数组记录该子树有的颜色并且子树外也有

#include<stdio.h>
#include<map>
#include<vector>
using namespace std;
const int N=1e5+5;
int rt[N],dep[N];
int a[N],b[N];
int sum[N];
int ans[N],n;
vector<int> edge[N];
map<int,int> m[N];
map<int,int> M;
map<int,int>::iterator it;

int merge(int x,int y)
{
	if(m[x].size()<m[y].size())
	    swap(x,y);
	for(it=m[y].begin();it!=m[y].end();it++)
	{
		if(m[x][it->first]==0)
		sum[x]++;
		m[x][it->first]+=it->second;
		if(m[x][it->first]==M[it->first])
		sum[x]--;
	} 
	return x;
}

void dfs(int x,int fa)
{
	it=m[x].begin();
	if(it->second==M[it->first])
	sum[x]=0;
	else
	sum[x]=1;
	dep[x]=dep[fa]+1;
	for(int i=0;i<edge[x].size();i++)
	if(edge[x][i]!=fa)
	dfs(edge[x][i],x);
	for(int i=0;i<edge[x].size();i++)
	{
		if(edge[x][i]!=fa)
		rt[x]=merge(rt[x],rt[edge[x][i]]);
	}
	ans[x]=sum[rt[x]];
}

int main()
{
	while(~scanf("%d",&n))
	{
		int c;
		for(int i=1;i<=n;i++)
		{
			scanf("%d",&c);
			m[i][c]++;
			M[c]++;
			rt[i]=i;
		}
		for(int i=1;i<n;i++)
		{
			scanf("%d%d",a+i,b+i);
			edge[a[i]].push_back(b[i]);
			edge[b[i]].push_back(a[i]);
		}
		dep[0]=0;
		dfs(1,0);
		for(int i=1;i<n;i++)
		{
			if(dep[a[i]]>dep[b[i]])
			printf("%d\n",ans[a[i]]);
			else
			printf("%d\n",ans[b[i]]);
		}
		
		for(int i=1;i<=n;i++)
		{
			edge[i].clear();
			m[i].clear();
		}
		M.clear();
	}
}