CF 246E Blood Cousins Return （dsu on tree）

題目連結：https://codeforces.com/contest/246/problem/E

題意：n個點，每個節點有一個名字，組成一個森林。m個詢問，每個詢問由（v，k）組成，問v的子樹中，深度為k（v的深度為1）的的節點有多少不同的名字。

思路：将詢問轉化為子樹中，詢問（v，k）問的就是v的子樹中，深度為deep（v）+k的節點有多少不同的名字。維護子樹中，每個深度不同名字的個數即可，我用的set。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 1e5+10;
int n,m;
char s[55];
unordered_map<string,int> mp;
int a[N];
bool book[N];
struct edge{
	int to,nex;
}g[N];
int head[N],cnt;
struct node{
	int id,dep,nex;
}q[N];
int head1[N],cnt1;
int sz[N],dep[N],son[N]; 
void getsz(int u){
	sz[u]=1;
	int mx=-1;
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		dep[v]=dep[u]+1;
		getsz(v);
		if(mx<sz[v]){
			mx=sz[v];
			son[u]=v;
		}
		sz[u]+=sz[v];
	}
}
int SON;
set<int> num[N];
void dfs1(int u,bool op){	
	if(op) num[dep[u]].insert(a[u]);
	else if(!num[dep[u]].empty())num[dep[u]].clear();
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		if(v==SON) continue;
		dfs1(v,op);
	}	
}
int ans[N];
void dfs(int u,bool keep){
	for(int i=head[u];i!=-1;i=g[i].nex){
		int v=g[i].to;
		if(v==son[u]) continue;
		dfs(v,0);
	}
	if(son[u]!=-1){
		dfs(son[u],1);
		SON=son[u];
	}
	dfs1(u,1);
	SON=0;
	for(int i=head1[u];~i;i=q[i].nex){
		int deep=q[i].dep;
		ans[q[i].id]=num[deep].size();
	}
	if(keep==0)
		dfs1(u,0); 
}
void add(int u,int v){
	g[cnt]=edge{v,head[u]};
	head[u]=cnt++;
}
void add1(int u,int id,int de){
	q[cnt1]=node{id,de,head1[u]};
	head1[u]=cnt1++;
}
int main(void){
	num[1].clear();
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		son[i]=head[i]=head1[i]=-1;
	for(int i=1;i<=n;i++){
		int u;
		scanf("%s%d",s+1,&u);
		if(!mp[s+1]) mp[s+1]=i;
		a[i]=mp[s+1];
		if(u) add(u,i);
		book[i]=(u==0);
	}
	for(int i=1;i<=n;i++)
		if(book[i]){
			dep[i]=1;
			getsz(i);
		} 
	scanf("%d",&m);
	for(int i=1;i<=m;i++){
		int u,k;
		scanf("%d%d",&u,&k);
		if(dep[u]+k>n) continue;
		add1(u,i,k+dep[u]);
	}
	for(int i=1;i<=n;i++)
		if(book[i]){
			dep[i]=1;
			dfs(i,0);
		} 
	for(int i=1;i<=m;i++)
		printf("%d\n",ans[i]);
	return 0;	
}