🐈

ABC202 E - Count Descendants

3 min read

ABC202 E - Count Descendants

様々な解法が存在する問題。
この記事では累積和の考えを利用したDFSを用いた解法を示す。

問題

https://atcoder.jp/contests/abc202/tasks/abc202_e

概要

N頂点の木と、Q個のクエリが与えられる。
木の根を頂点1とする。

クエリの内容

整数U_i,D_iが与えられるので、以下の条件を満たす頂点の個数を求めよ。

  • 根への最短パス上に頂点U_iが存在する
  • 根への最短パスまでの距離がD_iである

制約

2 \leq N \leq 2 \times 10 ^ 5 \\ 2 \leq Q \leq 2 \times 10 ^ 5

解法

問題は以下のように言い換えられる。

  • U_iを根とする部分木にある頂点のうち、深さがD_iである頂点の数を求めよ

各クエリごとに頂点の数を数えると、O(NQ)となり間に合わない。

そこで、クエリをオフラインで処理することを考える。
各クエリの答えの配列をresとする。

事前処理として、以下を行う。

  • 各頂点にクエリの情報を持たせる(クエリ番号, D_iの値)
  • 各頂点の深さを求める(これは後の探索と同時に行っても良い)

次に、頂点1からスタートする深さ優先探索をする。探索を行いながら、今まで訪問した頂点の深さをカウントした配列を更新する。この配列をcurとする。
深さ優先探索をしながら、各頂点に到達した時と、探索し終わった時に以下の操作を行う。

  • 到達した時
    各クエリに対して、res[クエリ番号] = -cur[D_i]とする。

  • 探索し終わった時
    各クエリに対して、res[クエリ番号] = res[クエリ番号] + cur[D_i]とする。

これで各クエリに対する答えがresに求まる。

計算量はO(N + Q)となり、十分間に合う。

解法の意味

上記の解法の意味を示す。
例えば以下のような木があるとする。

探索の訪問順は、以下のようになる。

1 -> 2 -> 4(訪問) -> 8 -> 4 -> 9 -> 10 -> 9 -> 4(探索終わり) -> 2 -> ...

ここでは、頂点4に対するクエリを答える場合を考える。
頂点4に訪問する直前でのcurの配列の中身は次のようになる。

その後、頂点4, 8, 9, 10を訪問した際にcurの中身を更新し、頂点4の探索が終わった時点でのcurの中身は次のようになる。

curの中身を見ると、頂点4を根とする部分木に含まれる頂点の深さの個数が差分として現れていることがわかる。

つまり、各頂点に対するクエリを答えるためには、各頂点を訪問する直前のcurの状態と、探索が終わった時点でのcurの状態がわかれば良いことがわかる。

しかし、curの状態を各頂点ごとに保持しようと考えると、各頂点ごとに大きさNの配列を持たなければいけないため、空間計算量は最大でO(N ^ 2)となってしまう。

そこで、ある頂点uのクエリに答えるには、uが持つ各クエリの値に対して、uに訪問した時点のcur[D_i]の値(つまりuに到達するまでに通った深さD_iの頂点の数)を引いておき、uの探索が終了した時点のcur[D_i]の値(つまりuに到達するまでに通った深さD_iの頂点の数と頂点uを根とする部分木に含まれる深さD_iの頂点の数の和)を足すことで、差分である頂点uを根とする部分木に含まれる深さD_iの頂点の数が求められることになる。

こうすることで、各頂点に対して、その頂点が持つクエリの分の情報だけを持つだけで十分になり、全体の空間計算量はO(Q)となる。

全体の時間計算量はO(N + Q)である。

提出コード

ACコードは以下。

https://atcoder.jp/contests/abc202/submissions/22853073
#include <bits/stdc++.h>
using namespace std;

const int N = 210000;

// 各クエリの結果
vector<int> res(N);
// クエリの情報
// first := クエリ番号
// second := クエリのDiの値
vector<vector<pair<int, int>>> id(N);
// 辺
vector<vector<int>> edges(N);
// 各頂点の深さ
vector<int> depth(N, 0);
// 探索時に更新する深さをカウントする配列
vector<int> cur(N, 0);

// 根からの探索
void dfs(int v = 0, int par = -1) {
    // 各クエリについて、訪問時にcur[Di]の値を引いておく
    for (int i = 0; i < id[v].size(); i++) {
        int idx = id[v][i].first, d = id[v][i].second;
        res[idx] -= cur[d];
    }

    // curの中身を更新する
    cur[depth[v]]++;

    // 子供への探索
    for (int u : edges[v]) {
        if (u == par) continue;
        depth[u] = depth[v] + 1;
        dfs(u, v);
    }

    // 各クエリについて、探索終了時にcur[Di]の値を足す
    for (int i = 0; i < id[v].size(); i++) {
        int idx = id[v][i].first, d = id[v][i].second;
        res[idx] += cur[d];
    }
}

int main() {
    // 入力受け取り
    int n;
    cin >> n;
    for (int i = 1; i < n; i++) {
        int p;
        cin >> p;
        p--;
        edges[p].push_back(i);
        edges[i].push_back(p);
    }

    int q;
    cin >> q;

    for (int i = 0; i < q; i++) {
        int u, d;
        cin >> u >> d;
        u--;
        id[u].push_back({i, d});
    }

    // 探索
    dfs();

    // 答えの出力
    for (int i = 0; i < q; i++) cout << res[i] << endl;
}

Discussion

ログインするとコメントできます