【CF246E】Blood Cousins Return【dsu on tree】

传送门

Solution

在\(dsu\ on\ tree\)时每个深度维护一个\(map\)保存这个深度上点的名字
查询时输出对应\(map\)的大小即可

#include<bits/stdc++.h>
using namespace std;
const int N=2e5+10;
int n,m,cnt,first[N],ans[N];
struct node{
	int v,nxt;
}e[N<<1];
vector<pair<int,int> > que[N];
inline void add(int u,int v){
	e[++cnt].v=v;e[cnt].nxt=first[u];first[u]=cnt;
}
int dep[N],hson[N],siz[N],pa[N][20];
string s[N];
inline void dfs(int u,int f){
	pa[u][0]=f;siz[u]=1;
	for(int i=1;i<=19;++i) pa[u][i]=pa[pa[u][i-1]][i-1];
	for(int i=first[u];i;i=e[i].nxt){
		int v=e[i].v;
		if(v==f) continue;
		dep[v]=dep[u]+1;
		dfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[hson[u]]) hson[u]=v;
	}
}
inline int up(int u,int k){
	for(int i=19;i>=0;--i)
		if(k&(1<<i)) u=pa[u][i];
	return u;
}
int son;
map<string,int> w[N]; 
inline void work(int u,int f,int tp){
	if(w[dep[u]].find(s[u])!=w[dep[u]].end()){
		w[dep[u]][s[u]]+=tp;
		if(!w[dep[u]][s[u]]) w[dep[u]].erase(w[dep[u]].find(s[u]));
	}
	else if(tp==1) w[dep[u]][s[u]]=1;
	for(int i=first[u];i;i=e[i].nxt){
		int v=e[i].v;
		if(v==f||v==son) continue;
		work(v,u,tp);
	}
}
inline void dsu(int u,int f,int tp){
	for(int i=first[u];i;i=e[i].nxt){
		int v=e[i].v;
		if(v==f||v==hson[u]) continue;
		dsu(v,u,0); 
	}
	if(hson[u]) dsu(hson[u],u,1),son=hson[u];
	work(u,f,1);son=0;
	for(int i=0;i<que[u].size();++i) ans[que[u][i].second]=w[que[u][i].first].size();
	if(!tp) work(u,f,-1);
}
char c[10];
int main(){
	scanf("%d",&n);
	for(int i=1,f;i<=n;++i){
		scanf("%s",c);	
		s[i]=c;
		scanf("%d",&f);
		if(!f) f=n+1;
		add(f,i);
	}
	scanf("%d",&m);
	dfs(n+1,0);
	for(int i=1,u,k;i<=m;++i){
		scanf("%d%d",&u,&k);
		que[u].push_back(make_pair(dep[u]+k,i));
	}
	dsu(n+1,0,1);
	for(int i=1;i<=m;++i) printf("%d\n",ans[i]);
	return 0;
}