CodeForces - 375D Tree and Queries(树上启发式合并)

题目链接：点击查看

题目大意：给出一棵有根树，每个节点都有一个编号代表颜色，现在给出m个询问，每个询问的形式为u k，需要回答以u为根节点的子树中，数量超过k的颜色有多少种

题目分析：树上启发式合并模板题，只不过一开始我不太会处理如何快速获取超过k的颜色有多少种，傻傻的想了半天数据结构无果，想用map二分，结果发现只能对key二分，而不能对val二分，如果模拟的话又太麻烦，如果直接暴力的话时间复杂度又上来了，想了半天无果后去看了看别人的代码，发现了一种真的很巧妙的方法，就是用另一个cntt数组，记录cnt数组每个出现次数的次数，换句话说每次当cnt数组新纪录颜色之后，用cntt数组记录一下cnt数组，会造成的影响是，若一个颜色出现了n次，那么cnt[color]=n，同时cntt[1]=cntt[2]=....=cntt[n-1]=cntt[n]=1，这样如果我们想直接查找出现k次的颜色，直接访问cntt[k]就是我们想要的结果了，想一想是不是

剩下的就没什么难度了，不过对于这个题目有个小坑，虽然题目说了给出的树是以点 1 为根的有根树，但题目给数据时给出的边却是以无向边的形式给出的，所以我们需要按照无向边的形式存边，访问的时候还是用v!=fa的形式访问就好了

代码：

#include<iostream>
#include<cstdio> 
#include<string>
#include<ctime>
#include<cstring>
#include<algorithm>
#include<stack>
#include<queue>
#include<map>
#include<set>
#include<sstream>
#include<unordered_map>
using namespace std;

typedef long long LL;

const int inf=0x3f3f3f3f;

const int N=1e5+100;

int color[N],ans[N],son[N],num[N],cnt[N],cntt[N];//cnt记录颜色数,cntt记录cnt的个数 

bool vis[N];

vector<int>node[N];

vector<pair<int,int>>Q[N];//<id,num>

void dfs_son(int u,int fa)
{
	num[u]=1;
	son[u]=-1;
	for(auto v:node[u])
	{
		if(v==fa)
			continue;
		dfs_son(v,u);
		num[u]+=num[v];
		if(son[u]==-1||num[v]>num[son[u]])
			son[u]=v;
	}
}

void cal(int u,int fa,int val)
{
	if(val==-1)
		cntt[cnt[color[u]]]--;
	cnt[color[u]]+=val;
	if(val==1)
		cntt[cnt[color[u]]]++;
	for(auto v:node[u])
	{
		if(vis[v]||v==fa)
			continue;
		cal(v,u,val);
	}
}

void dfs(int u,int fa,int keep)
{
	for(auto v:node[u])
	{
		if(v==son[u]||v==fa)
			continue;
		dfs(v,u,0);
	}
	if(son[u]!=-1)
	{
		dfs(son[u],u,1);
		vis[son[u]]=true;
	}
	cal(u,fa,1);
	for(auto i:Q[u])
	{
		int id=i.first;
		int k=i.second;
		ans[id]=cntt[k];
	}
	if(son[u]!=-1)
		vis[son[u]]=false;
	if(!keep)
		cal(u,fa,-1);
}

int main()
{
//	freopen("input.txt","r",stdin);
//	ios::sync_with_stdio(false);
	int n,m;
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++)
		scanf("%d",color+i);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		node[u].push_back(v);
		node[v].push_back(u);
	}
	for(int i=1;i<=m;i++)
	{
		int u,k;
		scanf("%d%d",&u,&k);
		Q[u].push_back(make_pair(i,k));
	}
	dfs_son(1,-1);
	dfs(1,-1,1);
	for(int i=1;i<=m;i++)
		printf("%d\n",ans[i]);
	
	
	
	
	
	
	
	
	
	return 0;
}