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: