Blood Cousins (dsu on tree + 求第k级祖先)

题目链接: Blood Cousins

大致题意

给出一片森林, 询问对于给定的点而言, 其k级祖先中有多少个点的深度和当前节点相同.

解题思路

dsu on tree

对于每个询问x, k而言, 我们假设x的k级祖先为p, 则对于每一个p, 我们求出其子树中不同深度的节点数量即可. 我们可以通过dsu on tree做法求出.

关于节点的k级祖先求法, 大家可以看这篇博客 ➡️ 博客链接

AC代码

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 1; i <= (n); ++i)
using namespace std;
typedef long long ll;
const int N = 1E5 + 10;
vector<int> edge[N];

int son[N], sz[N];
int depth[N], f[N][18];
void dfs1(int x, int fa) {
	depth[x] = depth[fa] + 1;
	sz[x] = 1;

	f[x][0] = fa;
	rep(i, 17) f[x][i] = f[f[x][i - 1]][i - 1];

	for (auto& to : edge[x]) {
		if (to == fa) continue;
		dfs1(to, x);
		sz[x] += sz[to];
		if (sz[to] > sz[son[x]]) son[x] = to;
	}
}

int getk(int x, int k) {
	if (!k) return x;
	int index = log2(k);
	return getk(f[x][index], k - (1 << index));
}

vector<pair<int, int>> query[N]; // val, id
int res[N];
unordered_map<int, int> mp; // depth, cou;
void calc(int x, int fa, int pson) {
	mp[depth[x]]++;
	for (auto& to : edge[x]) {
		if (to == fa or to == pson) continue;
		calc(to, x, pson);
	}
}
void dfs2(int x, int fa, int tp) {
	for (auto& to : edge[x]) {
		if (to == fa or to == son[x]) continue;
		dfs2(to, x, 0);
	}

	if (son[x]) dfs2(son[x], x, 1);
	calc(x, fa, son[x]);

	for (auto& [val, id] : query[x]) res[id] = mp[val] - 1;

	if (!tp) mp.clear();
}
int main()
{
	int n; cin >> n;
	rep(i, n) {
		int p; scanf("%d", &p);
		edge[p].push_back(i);
	}
	for (auto& op : edge[0]) dfs1(op, 0);

	int m; cin >> m;
	rep(i, m) {
		int x, k; scanf("%d %d", &x, &k);
		int p = getk(x, k);
		query[p].push_back({ k + depth[p], i }); //采用相对于根节点的深度.
	}

	for (auto& op : edge[0]) dfs2(op, 0, 0); //由于是森林, 所以根节点应为轻儿子, 否则WA9

	rep(i, m) printf("%d%c", res[i], " \n"[i == m]);
	return 0;
}

END