Problem - 208E - Codeforces （树上启发式合并+倍增）

Problem - 208E - Codeforces

题意

给定一棵描述 $n$ 个人的家族关系树，两个人有相同的 $k$ 祖先，则两个人就是 $k$ 兄弟。给出 $m$ 次询问，询问 $v$ 有几个 $k$ 兄弟。

解法

可以先用倍增找到 $v$ 的 $k$ 祖先 $r$，那么问题就转换成了求 $r$ 的某深度的孩子的数量。这就变成了树上启发式合并的板子题了。

解法

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e5 + 5;
int n, q;
vector<int> g[maxn];
struct node {
    int id, k;
};
vector<node> Q[maxn];
int fa[maxn][21], sz[maxn], son[maxn], dep[maxn];
int res[maxn], cnt[maxn], flag;
void dfs1(int u, int f) {
    sz[u] = 1, dep[u] = dep[f] + 1;
    for (int i = 1; (1 << i) <= dep[u]; ++i)
        fa[u][i] = fa[fa[u][i - 1]][i - 1];
    for (auto v : g[u]) {
        if (v == f) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if (sz[v] > sz[son[u]]) son[u] = v;
    }
}
void count(int u, int f, int val) {
    cnt[dep[u]] += val;
    for (auto v : g[u]) {
        if (v == f || v == flag) continue;
        count(v, u, val);
    }
}
void query(int u) {
    for (auto e : Q[u]) {
        int d = dep[u] + e.k;
        res[e.id] = cnt[d];
    }
}
void dfs(int u, int f, bool keep) {
    for (auto v : g[u]) {
        if (v == f || v == son[u]) continue;
        dfs(v, u, 0);
    }
    if (son[u]) {
        dfs(son[u], u, 1);
        flag = son[u];
    }
    count(u, f, 1);
    query(u);
    flag = 0;
    if (!keep) count(u, f, -1);
}
int qfa(int u, int k) {
    if (dep[u] - k < 1) return 0;
    for (int i = 20; i >= 0; --i) {
        if (k & (1 << i)) u = fa[u][i];
    }
    return u;
}
int main() {
    scanf("%d", &n);
    rep(i, 1, n) {
        scanf("%d", &fa[i][0]);
        if (fa[i][0]) g[fa[i][0]].push_back(i);
    }
    rep(i, 1, n) if (!fa[i][0]) dfs1(i, 0);
    scanf("%d", &q);
    rep(i, 1, q) {
        int u, k;
        scanf("%d%d", &u, &k);
        u = qfa(u, k);
        Q[u].push_back({i, k});
    }
    rep(i, 1, n) if (!fa[i][0]) dfs(i, 0, 0);
    rep(i, 1, q) printf("%d%c", max(0, res[i] - 1), " \n"[i == q]);
}