dsu on tree(264E BLOOD COUSIN RETURNS), C++ (gcc)

dsu on tree(264E BLOOD COUSIN RETURNS)

#include <bits/stdc++.h>

using namespace std;
#define pb push_back
#define sz(v) (int)v.size()

typedef long long int ll;

struct query{
    int i, h;
    
    query(int ind, int height) : i(ind), h(height){}
};

const int N = 1e5 + 1;

vector<int> g[N];

ll big[N] = {0}, sub[N], col[N], total[2*N] = {0}, ans[N] = {0};

vector<query> q[N];

void dfs(int s, int p){
    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]);
        }
}

map<int, int> cnt_[N];

void add(int s, int p, int v, int lvl){
    int cur = col[s];
    int &cnt = cnt_[cur][lvl];
    cnt += v;
    if(cnt == 1 && v == 1) total[lvl]++;
    else if(cnt == 0 && v == -1) total[lvl]--;
    
    for(auto i : g[s])
        if(i != p && !big[i]) add(i, s, v, lvl+1);
}

void dfs_dsu(int s, int p, bool keep, int lvl){
    int x = g[s][0];
    //cout << s << " " << p << " " << lvl << " " << x << endl;
    for(auto i : g[s])
        if(i != x && i != p) dfs_dsu(i, s, 0, lvl+1);
    if(x != p) dfs_dsu(x, s, 1, lvl+1), big[x] = 1;
        
    add(s, p, 1, lvl);
    
    for(auto i : q[s]){
        const int &h = i.h;
        ans[i.i] = total[lvl+h];
    }
    
    if(x != p) big[x] = 0;
    if(!keep) add(s, p, -1, lvl);
}

map<string, int> m;
int ind = 1;

int main(){
    ios_base :: sync_with_stdio(0);
    cin.tie(0);
    
    vector<int> roots;
    int n;
    cin >> n;
    for(int j = 1; j <= n; j++){
        string s;
        cin >> s;
        int &cur = m[s];
        if(cur == 0) cur = ind, col[j] = ind, ind++;
        else col[j] = cur;
        int p;
        cin >> p;
        if(p == 0) roots.pb(j), g[j].pb(0);
        else g[j].pb(p), g[p].pb(j);
    }
    int m;
    cin >> m;
    for(int j = 0; j < m; j++){
        int u, h;
        cin >> u >> h;
        q[u].pb(query(j, h));
    }
    for(auto j : roots){
        dfs(j, 0);
        dfs_dsu(j, 0, 0, 1);
    }
    for(int j = 0; j < m; j++)
        cout << ans[j] << endl;
    //cout << endl;
    return 0;
}


run \| edit \| history \| help	0

λ

                                      .NET NoSQL database for rapid development