E - Count Descendants Editorial by TumoiYorozu
クエリ先読み + DFS による O(N+Q) 解法
問題は、ある頂点 \(U_i\) かその子孫の中(部分木)で、深さが \(D_i\) の頂点の個数を求める問題と言い換えることが出来ます。
まず、クエリを全て先読みして、各頂点に求めたい子のDepthを記録します(コードのQd
に対応)。
このとき、クエリが何番目のクエリだったかも同時に記録しておくと、復元にソートやmapなどが必要なくなり、\(O(Q)\) で復元可能になります(コードのQi
に対応)。
さて、根から順にDFSを行うわけですが、今までに通ったそれぞれの depth の個数を記録します(コードのdepth_count
に対応)。
dfsである頂点に注目したとき、子のdfsをする前にそれぞれのクエリに対し、元々のdepthが何個だったかを記録し(コードのDin
)、子のdfsをしたあとに、それぞれのdepthが何個増えたかの差分を取ることで、高速にカウントが可能です。
コード (提出)
#include <bits/stdc++.h>
using namespace std;
#define rep(i,a) for(int i=0;i<(a);i++)
int res[200005];
vector<int> child[200005];
vector<int> Qd[200005];
vector<int> Qi[200005];
vector<int> Din[200005];
int depth_count[200005] = {};
void dfs(int p, int depth) {
int Qn = Qd[p].size();
Din[p].resize(Qn);
rep(i, Qn) {
Din[p][i] = depth_count[Qd[p][i]];
}
depth_count[depth]++;
for(auto c : child[p]) {
dfs(c, depth + 1);
}
rep(i, Qn) {
int r = depth_count[Qd[p][i]] - Din[p][i];
res[Qi[p][i]] = r;
}
}
int main(){
int n;
scanf("%d", &n);
rep(i, n-1) {
int p;
scanf("%d", &p);
child[p-1].push_back(i+1);
}
int q;
scanf("%d", &q);
rep(i, q) {
int u, d;
scanf("%d%d", &u, &d);
Qd[u-1].push_back(d);
Qi[u-1].push_back(i);
}
dfs(0, 0);
rep(i, q) {
printf("%d\n", res[i]);
}
return 0;
}
posted:
last update: