Codeforces208E Blood Cousins

dsu on tree

題目連接

點我跳轉

題目大意

給你一片森林，每次詢問一個點與多少個點擁有共同的\(K\)級祖先

解題思路

詢問一個點有多少個共同的\(K\)級祖先

能夠等價於問它的\(K\)級祖先有多少個深度爲 \(dep + K\) 的子節點 (\(dep\)爲\(K\)級祖先的深度)

那麼對於每一個詢問要先倍增或者樹鏈剖分求出它的\(K\)級祖先，再將詢問保存在\(K\)級祖先的\(vector\)裏

而後跑一遍 dsu on tree 統計 + 保存答案，最後再將保存的答案按照讀入的詢問順序排個序輸出便可

AC_Code

#include<bits/stdc++.h>
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define per(i,n,a) for (int i=n;i>=a;i--)
#define pb push_back
#define fi first
#define se second
using namespace std;
const int N = 5e5 + 10;
struct Edge{
	int nex , to;
}edge[N << 1];
int head[N] , tot;
int a[N] , dep[N] , f[N][30] , siz[N] , hson[N] , HH;
vector<pair<int , int>>Q[N] , ans;
unordered_map<int , int>cnt;
void add_edge(int u , int v)
{
	edge[++ tot].nex = head[u];
	edge[tot].to = v;
	head[u] = tot;
}
void dfs(int u , int far)
{
	siz[u] = 1;
	dep[u] = dep[far] + 1;
	f[u][0] = far;
	for(int i = 1 ; (1 << i) <= dep[u] ; i ++)	
		f[u][i] = f[f[u][i - 1]][i - 1];
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far) continue ;		
		dfs(v , u);
		siz[u] += siz[v];
		if(siz[v] > siz[hson[u]]) hson[u] = v;
 	}
} 
void calc(int u , int far , int val , int dep)
{
	cnt[dep] += val;
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == HH) continue ;
		calc(v , u , val , dep + 1);
	}
}
void dsu(int u , int far, int op , int dep)
{
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == hson[u]) continue ;
		dsu(v , u , 0 , dep + 1);
	}
	if(hson[u]) dsu(hson[u] , u , 1 , dep + 1) , HH = hson[u];
	calc(u , far , 1 , dep);
	for(auto i : Q[u])
	{
		int k = i.fi , id = i.se;
		ans.pb(make_pair(id , cnt[dep + k] - 1));
	}
	HH = 0;
	if(!op) calc(u , far , -1 , dep);
}
int get_far(int x , int k)
{
	if(!k) return x;
    int t = dep[x] - k;
    per(i , 20 , 0) if(dep[f[x][i]] > t) x = f[x][i];
    return f[x][0];
}
signed main()
{
	int n , q;
	cin >> n;
	rep(i , 1 , n)
	{
		cin >> a[i];
		if(!a[i]) continue ;
		add_edge(i , a[i]) , add_edge(a[i] , i);
	}
	rep(i , 1 , n)
	{
		if(!a[i]) dfs(i , 0);
	}
	cin >> q;
	rep(i , 1 , q)
	{
		int a , b;
		cin >> a >> b;
		int x = get_far(a , b);
		if(!x) {ans.pb(make_pair(i , 0)) ; continue ;}
		Q[x].pb(make_pair(b , i));
	}
	rep(i , 1 , n)
	{
		if(!a[i]) dsu(i , 0 , 0 , 1) , cnt.clear();
	}
	sort(ans.begin() , ans.end());
	for(auto i : ans) cout << i.se << " ";
	cout << '\n';
	return 0;
}