🐈

ABC202 E - Count Descendants

2021/05/23に公開

C++

競技プログラミング

AtCoder

tech

ABC202 E - Count Descendants

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

問題

概要

$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コードは以下。

#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;
}

ABC202 E - Count Descendants

問題

概要

クエリの内容

制約

解法

解法の意味

提出コード

Discussion