E. Blood Cousins (k级祖先+树上启发式合并)

传送门

题意：给了一个森林，m次询问，每次问x,k，有多少个节点和x有相同的k级祖先。

首先这个询问是可以离线的，我们尝试把它转化成子树查询：

• 求出x的k级祖先记为an,于是查询就转化成了在an为根的子树中，找出有多少个节点和x的深度相同。
• 转化后的问题我们先考虑暴力求解:dfs，对于每个子树，暴力求出cnt[i]:这棵子树中，深度为i的节点有多少个，然后更新相应的查询。 O（n²）
• 优化：每递归完一棵子树，需要删除这颗子树的信息，但发现可以保留一棵子树的信息给根节点，可以树上启发式把时间优化成nlogn。

细节：可以用0节点把所有的树连起来，从而形成一棵树，只不过所有树的信息在递归完后需要擦除。

#include<bits/stdc++.h>
using namespace std;
//#pragma GCC optimize(2)
#define ull unsigned long long
#define ll long long
#define pii pair<int, int>
const int maxn = 1e5 + 10;
const ll mod = 998244353;
const ll inf = (ll)4e17+5;
const int INF = 1e9 + 7;
const double pi = acos(-1.0);
ll inv(ll b){if(b==1)return 1;return(mod-mod/b)*inv(mod%b)%mod;}
//给一个森林 m次询问 每次询问x k ,求有多少个点和x存在相同的k级祖先 
//0节点把森林变成一棵树 注意不同的树之间不能产生影响
//O(n)²的暴力方法:离线把查询转化成子树查询  用cnt[i]表示深度为i的节点有多少个
//dsu on tree 优化成nlogn
vector<int> g[maxn];
int fa[maxn];//父节点
int n,m;
vector<pii> q[maxn];//查询
int f[maxn][25];//i的2^j级祖先
int dep[maxn];
int size[maxn],son[maxn];//树剖求出重儿子
void dfs(int rt,int fa)//树上倍增求k级祖先   顺便求出重儿子
{	
	dep[rt]=dep[fa]+1;
	f[rt][0]=fa;
	size[rt]=1;
	for(int i=1;i<=20;i++) 
	{
		f[rt][i]=f[f[rt][i-1]][i-1];
	}
	for(int &i:g[rt]) 
	{
		dfs(i,rt);
		if(size[i] > size[son[rt]]) son[rt]=i;
		size[rt]+=size[i];
	}
}
int get_k(int x,int k)//x的k级祖先 不存在则返回0
{
	int t=dep[x]-k;
    for(int i=20;i>=0;i--)//倍增向上跳跃，求祖先
        if(dep[f[x][i]]>t) x=f[x][i];
    return f[x][0];
}
int cnt[maxn];
int ans[maxn];
int SON;
void add(int rt,int v) 
{
	cnt[dep[rt]]+=v;
	for(int i:g[rt]) 
	{
		if(i==SON) continue;
		add(i,v);
	}
}
void dfs2(int rt,bool save) 
{
	for(int i:g[rt]) 
	{
		if(i==son[rt]) continue;
		dfs2(i,0);
	}
	if(son[rt]) dfs2(son[rt],1),SON=son[rt];
	add(rt,1),SON=0;
	for(auto i:q[rt]) 
	{
		ans[i.second]=cnt[dep[rt]+i.first]-1;
	}
	if(!save) add(rt,-1);
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
	{
		scanf("%d",fa+i);
		g[fa[i]].push_back(i);
	}
	for(int i:g[0]) dfs(i,0);
	scanf("%d",&m);
	for(int i=1;i<=m;i++)
	{
		int x,k;
		scanf("%d %d",&x,&k);
		int an=get_k(x,k);
		q[an].push_back({k,i});//查询an子树中  dep[an]+k层的节点有多少个
	}
	for(int i:g[0]) dfs2(i,0);//注意传0 因为树与树之间不会产生影响
	for(int i=1;i<=m;i++) printf("%d ",ans[i]);
	return 0;
}