CodeForces - 208E LCA+树上启发式合并

题目链接：https://codeforces.com/problemset/problem/208/E
题目大意：给你一片森林，再给q次询问：每次x y：询问x的有级祖先有多少个儿子和x处于同一层(排除x节点外)。

思路：我们把查询改为x的k级祖先y在deep[x]层有多少个子节点。用LCA处理一下，然后启发式合并可以了。

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

vector<int> G[100005];
int c[100005], ans[100005], vis[100005];
struct LCA{
    int d[100005], fa[100005][22], lg[100005];
    void init(int n, int root){//预处理
        if(!lg[10]){
            for(int i = 1; i <= n; ++i){
                lg[i] = lg[i-1] + (1 << lg[i-1] == i);
            }
        }
        dfs(root, 0);
    }
    void dfs(int now, int father){
        vis[now]=1;
        fa[now][0]=father; d[now]=d[father]+1;
        for(int i=1; i<=lg[d[now]]; ++i){
            fa[now][i]=fa[fa[now][i-1]][i-1];
        }
        for(auto x: G[now]){
            if(x!=father){
                dfs(x, now);
            }
        }
    }
    int klca(int x, int y){//第k级祖先(自己第0级)
        for(int i=30; i>=0; i--){
            if(y&(1<<i)){
                x=fa[x][i];
            }
        }
        return x;
    }
}lca;

vector<pair<int, int> > q[100005];
int siz[100005], dfn[100005], son[100005], d[100005], T=0;
void dfs(int u, int deep){
    vis[u]=1;
    siz[u]=1, dfn[u]=++T; d[T]=deep;
    for(auto x: G[u]){
        dfs(x, deep+1);
        siz[u]+=siz[x];
        son[u]=siz[x]>siz[son[u]]?x:son[u];
    }
}

void DSU(int u, int k){
    vis[u]=1;
    for(auto x: G[u]){
        if(x!=son[u]){
            DSU(x, 0);
        }
    }
    if(son[u]) DSU(son[u], 1);
    for(auto x: G[u]){
        if(x!=son[u]){
            for(int i=0; i<siz[x]; i++){
                int now=dfn[x]+i;
                c[d[now]]++;
            }
        }
    }
    c[d[dfn[u]]]++;
    for(auto x: q[u]){
        int k=x.first, pos=x.second;
        ans[pos]=c[k+d[dfn[u]]];
    }
    if(!k){

        for(int i=0; i<siz[u]; ++i){
            int now=dfn[u]+i;
            c[d[now]]=0;
        }
    }
}

int main(){
    int n, x, k; scanf("%d", &n);
    for(int i=1; i<=n; i++){
        scanf("%d", &x);
        if(x){
            G[x].push_back(i);
        }
    }
    for(int i=1; i<=n; i++){
        if(!vis[i]){
            lca.init(n, i);
        }
    }
    memset(vis, 0, sizeof(vis));
    int Q; scanf("%d", &Q);
    for(int i=1; i<=Q; i++){
        scanf("%d%d", &x, &k);
        int pos=lca.klca(x, k);
        q[pos].push_back({k, i});
    }
    for(int i=1; i<=n; i++){
        if(!vis[i]){
            dfs(i, 0);
        }
    }
    memset(vis, 0, sizeof(vis));
    for(int i=1; i<=n; i++){
        if(!vis[i]){
            DSU(i, 0);
//            for(int i=1; i<=5; i++){
//                cout<<i<<" "<<c[i]<<endl;
//            }
        }
    }
    for(int i=1; i<=Q; i++){
        printf("%d ", (ans[i]?ans[i]-1:0));
    }
    printf("\n");

    return 0;
}
/*
6
0 1 1 0 4 4
1
5 1
*/