CodeForces 246E.Blood Cousins Return【树上启发式合并】

传送门

题意

给定一片森林，每次询问一个节点的K-Son共有个多少不同的名字。一个节点的K-Son即为深度是该节点深度加K的节点。

题解

和前一道题基本是同一个套路，只是这个有重名，所以直接用set数组来代替int数组计数就行了。
因为重儿子没有求对吃了一大波TLE。

代码

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <queue>
#include <map>
#include <set>
#define xx first
#define yy second
using namespace std;
typedef long long LL;
typedef pair<int,int> PII;
const int inf=0x3f3f3f3f;
const LL INF=0x3f3f3f3f3f3f3f3f;
const int N=3e5+10;
const int M=1e6+10;
int n,dep[N],siz[N],son[N],ans[N],m;
int s[N];
int head[N],to[N*2],nxt[N*2],tot;
void add(int u,int v){to[++tot]=v;nxt[tot]=head[u];head[u]=tot;}
vector<PII> q[N];
map<string,int> mp;
set<int> st[N];

void predfs(int u,int rt){
	siz[u]=1;dep[u]=dep[rt]+1;
	for(int i=head[u];i;i=nxt[i]){
		int v=to[i];
		if(v==rt) continue;
		predfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

void upd(int u,int fa,bool opt){
	if(opt) st[dep[u]].insert(s[u]);
	else st[dep[u]].erase(s[u]);
	for(int i=head[u];i;i=nxt[i]) if(to[i]!=fa) upd(to[i],u,opt);
}

void dfs(int u,int fa,bool keep){
	for(int i=head[u];i;i=nxt[i]){
		if(to[i]==fa||to[i]==son[u]) continue;
		dfs(to[i],u,false);
	}
	if(son[u]) dfs(son[u],u,true);
	st[dep[u]].insert(s[u]);
	for(int i=head[u];i;i=nxt[i]) if(to[i]!=fa&&to[i]!=son[u]) upd(to[i],u,true);
	for(PII iq:q[u]) if(iq.xx<=n) ans[iq.yy]=st[iq.xx].size();
	if(!keep) upd(u,fa,false);
	
}

int main(){
	ios::sync_with_stdio(false);cin.tie(0);
	cin>>n;
	for(int i=1,u;i<=n;i++){
		string temp;
		cin>>temp>>u;
		if(!mp[temp]) mp[temp]=mp.size();
		s[i]=mp[temp];
		add(u,i);add(i,u);
	}
	for(int i=head[0];i;i=nxt[i]) predfs(to[i],0);
	cin>>m;
	for(int i=1,u,w;i<=m;i++){
		cin>>u>>w;
		q[u].push_back({dep[u]+w,i});
	}
	for(int i=head[0];i;i=nxt[i]) dfs(to[i],0,false);
	for(int i=1;i<=m;i++) printf("%d
",ans[i]);
	return 0;
}