U41492 树上数颜色 题解

U41492 树上数颜色

题目描述

给一棵根为1的树，每次询问子树颜色种类数

输入格式

第一行一个整数n，表示树的结点数

接下来n-1行，每行一条边

接下来一行n个数，表示每个结点的颜色c[i]

接下来一个数m，表示询问数

接下来m行表示询问的子树

输出格式

对于每个询问，输出该子树颜色数

输入输出样例

输入 #1复制

5

1 2

1 3

2 4

2 5

1 2 2 3 3

5

1

2

3

4

5

输出 #1复制

3

2

1

1

1

说明/提示

对于前三组数据，有1≤m,c[i]≤n≤100

而对于所有数据，有1≤m,c[i]≤n≤1e5

暴力算法 ac

#include<iostream>
#include<vector>
using namespace std;
const int maxn=1e5+10;
vector<int> tr[maxn];
int tong[maxn],ans[maxn],c[maxn],fa[maxn],now;
void calc(int u,int f,int fl)
{
	tong[c[u]]+=fl;
	if (fl==1&&tong[c[u]]==1) now++;
	if (fl==-1&&tong[c[u]]==0) now--;
	for (int i=0;i<tr[u].size();i++)
	{
		int v=tr[u][i];
		if (v==f) continue;
		calc(v,u,fl);
	}
}
void dfs(int u,int f)
{
	for (int i=0;i<tr[u].size();i++)
	{
		int v=tr[u][i];
		if (v==f) continue;
		dfs(v,u);
	}
	calc(u,f,1);
	ans[u]=now;
	calc(u,f,-1);
}
int main()
{
	int n,m;
	cin>>n;
	for (int i=1;i<n;i++)
	{
		int x,y;
		cin>>x>>y;
		tr[x].push_back(y);
		tr[y].push_back(x);
	}
	for (int i=1;i<=n;i++) cin>>c[i];
	dfs(1,0);
	cin>>m;
	while(m--)
	{
		int x;
		cin>>x;
		cout<<ans[x]<<endl;	 
	} 
}