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

传送门

题意

给一片森林，(m) 次查询，询问每个节点与其它多少个节点有共同的第 (k) 祖先。

题解

对于这种无修改的询问题目，直接离线询问。
把每个询问中那个第 (k) 祖先用倍增求出来，然后问题转化为问第 (k) 祖先的子树中有多少个点与询问的点同深度。这个问题可以用树上启发式合并来解决。
其实 dsu on tree 的写法也比较灵活，完全可以不用照着模板改。

代码

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <queue>
#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,fa[N][22],dep[N],son[N],siz[N],m,ans[N],cnt[N];
vector<int> g[N];
vector<PII> ask[N];

void predfs(int u,int rt){
	fa[u][0]=rt;dep[u]=dep[rt]+1;siz[u]=1;
	for(int v:g[u]){
		if(v==rt) continue;
		predfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

void add(int u,int fa,int val){
	cnt[dep[u]]+=val;
	for(int v:g[u]) if(v!=fa) add(v,u,val);
}

void dfs(int u,int fa,bool keep){
	for(int v:g[u]){
		if(v==fa||v==son[u]) continue;
		dfs(v,u,false);
	}
	if(son[u]) dfs(son[u],u,true);
	cnt[dep[u]]++;
	for(int v:g[u]) if(v!=fa&&v!=son[u]) add(v,u,1);
	for(PII q:ask[u]) ans[q.yy]=cnt[q.xx]-1;
	if(!keep) add(u,fa,-1);
}

int main(){
	scanf("%d",&n);
	for(int i=1,u;i<=n;i++){
		scanf("%d",&u);
		g[u].push_back(i);
		g[i].push_back(u);
	}
	for(int i=1;i<=n;i++) if(!fa[i][0]) predfs(i,0);
	for(int j=1;j<=20;j++) for(int i=1;i<=n;i++) fa[i][j]=fa[fa[i][j-1]][j-1];
	scanf("%d",&m);
	for(int i=1,u,k;i<=m;i++){
		scanf("%d%d",&u,&k);
		int rt=u;
		for(int j=20;j>=0;j--) if(k>=(1<<j)) rt=fa[rt][j],k-=(1<<j);
		if(rt==0) ans[i]=0;
		ask[rt].push_back({dep[u],i});
	}
	for(int i=1;i<=n;i++) if(!fa[i][0]) dfs(i,0,false);
	for(int i=1;i<=m;i++) printf("%d ",ans[i]);
	return 0;
}