[CF208E] Blood Cousins - dsu on tree, LCA

Description

给你一片森林，每次询问一个点与多少个点拥有共同的 K 级祖先。

Solution

dsu on tree

首先，求出一个点的 k 级祖先，并且把询问挂在他身上，问的就是这个点 p 子树内有多少个深度为 d 的点

在 dfs 过程中，对每个点，递归处理各个孩子，同时保留其重孩子的贡献，再加上轻孩子的贡献，得到自身的贡献

总结一下我们需要支持哪些事情：求 k 级祖先（倍增处理），将询问挂在祖先上，树上启发式合并处理所有询问

#include <bits/stdc++.h>
using namespace std;

#define int long long

const int N = 1e5 + 5;

///////////////////////////////////////////////////

vector<int> g[N];
int n, wson[N], dep[N], siz[N], fa[N][20], ans[N];

///////////////////////////////////////////////////

vector<pair<int, int>> question[N];

///////////////////////////////////////////////////

int bucket[N];
vector<int> record;

void add(int p)
{
    bucket[p]++;
    record.push_back(p);
}

void dec(int p)
{
    bucket[p]--;
}

void clr()
{
    for (auto i : record)
        dec(i);
    record.clear();
}

//////////////////////////////////////////////////

int get_kth_father(int p, int k)
{
    for (int i = 19; i >= 0; i--)
        if ((k >> i) & 1)
            p = fa[p][i];
    return p;
}

void make_question(int p, int k, int id)
{
    int par = get_kth_father(p, k);
    if (par)
    {
        question[par].push_back(make_pair(dep[p], id));
    }
}

void dfs_pre(int p)
{
    siz[p] = 1;
    for (int q : g[p])
    {
        dep[q] = dep[p] + 1;
        fa[q][0] = p;
        for (int j = 1; j < 20; j++)
            fa[q][j] = fa[fa[q][j - 1]][j - 1];
        dfs_pre(q);
        siz[p] += siz[q];
        if (siz[wson[p]] < siz[q])
            wson[p] = q;
    }
}

void dfs_add(int p)
{
    add(dep[p]);
    for (int q : g[p])
        dfs_add(q);
}

void dfs(int p)
{
    for (int q : g[p])
    {
        if (q == wson[p])
            continue;
        dfs(q);
        clr();
    }
    if (wson[p])
    {
        dfs(wson[p]);
    }
    for (int q : g[p])
    {
        if (q == wson[p])
            continue;
        dfs_add(q);
    }
    add(dep[p]);

    // cout << "p=" << p << endl;
    // for (int i = 0; i <= 10; i++)
    //     cout << bucket[i] << " ";
    // cout << endl;
    for (auto [depth, id] : question[p])
    {
        // cout << " question: " << depth << " " << id << endl;
        ans[id] = bucket[depth] - 1;
    }
}

///////////////////////////////////////////////////

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n;
    for (int i = 1; i <= n; i++)
    {
        int r;
        cin >> r;
        g[r].push_back(i);
    }

    dfs_pre(0);

    int m;
    cin >> m;
    for (int i = 1; i <= m; i++)
    {
        int v, k;
        cin >> v >> k;
        make_question(v, k, i);
    }

    dfs(0);

    for (int i = 1; i <= m; i++)
        cout << ans[i] << " ";
}