dsu on tree (http://codeforces.com/contest/208/problem/E), C++ (gcc)

dsu on tree (http://codeforces.com/contest/208/problem/E)

#include <bits/stdc++.h>

using namespace std;
#define pii pair<int, int>
#define mp make_pair
#define pb push_back
#define sz(v) (int)v.size()
typedef long long int ll;

const int N = 1e5 + 5;

vector<int> g[N];

int sub[N], lvl[N], dp[20][N], cnt[N] = {0}, ans[N] = {0};

struct node{
    int i, v;
    
    node(int ind, int ver) : i(ind), v(ver){}
};

vector<node> q[N];

void dfs(int s, int p){
    lvl[s] = lvl[p] + 1;
    dp[0][s] = p;
    for(int j = 1; j < 20; j++) dp[j][s] = dp[j-1][dp[j-1][s]];
    sub[s] = 1;
    
    for(auto &i : g[s]) 
        if(i != p){
            dfs(i, s), sub[s] += sub[i];
            if(sub[i] > sub[g[s][0]]) swap(i, g[s][0]);
        }
}

bool big[N] = {0};

void add(int s, int p, int val){
    cnt[lvl[s]] += val;
    for(auto i : g[s])
        if(i != p && big[i] != 1) add(i, s, val);
}

inline void answer_query(int s){
    for(auto i : q[s])
        ans[i.i] = cnt[lvl[i.v]] - 1;
}

void dfs_HLD(int s, int p, bool keep){
    if(sz(g[s]) == 0) return;
    
    int bigchild = g[s][0];
    for(auto i : g[s])
        if(i != p && i != bigchild) dfs_HLD(i, s, 0);
    
    if(bigchild != p) big[bigchild] = 1, dfs_HLD(bigchild, s, 1);
    
    add(s, p, 1);
    answer_query(s);
    
    if(bigchild != p) big[bigchild] = 0;
    if(!keep) add(s, p, -1);
}

int node_at(int u, int raise){
    if(lvl[u] <= raise) return 0;
    for(int j = 19; j >= 0; j--)
        if(raise >= (1 << j)) raise -= (1 << j), u = dp[j][u];
    return u;
}

int main(){
    ios_base :: sync_with_stdio(0);
    cin.tie(0);
    
    memset(dp, 0, sizeof(dp));
    lvl[0] = 0;
    
    int n;
    vector<int> roots;
    cin >> n;
    for(int j = 1; j <= n; j++){
        int p;
        cin >> p;
        if(p) g[j].pb(p), g[p].pb(j);
        else roots.pb(j);
    }
    for(auto i : roots) dfs(i, 0);
    
    int query;
    cin >> query;
    for(int j = 0; j < query; j++){
        int v, p;
        cin >> v >> p;
        int raised = node_at(v, p);
        if(raised) q[raised].pb(node(j, v));
    }
    for(auto i : roots)
        dfs_HLD(i, 0, 0);
    
    for(int j = 0; j < query; j++) cout << ans[j] << " ";
    return 0;
}


run \| edit \| history \| help	0