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];
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;
return 0;
}
|
run
| edit
| history
| help
|
0
|
|
|