【ABC202】E - Count Descendants

atcoder.jp

公式解説とは違う方法らしい。

解法

この問題は、頂点 $ U_i $の子要素で根からの距離が $ D_i $である頂点はいくつか、というクエリに回答せよ、と言い換えられる。

複数のクエリに回答する問題

複数のクエリに回答する問題は、次のような解法パターンがある。

• それぞれのクエリで何かしらの調査をして答えを計算する
• 事前調査の結果を用いて、一気にクエリに回答する
• クエリを先読みしておき、調査しながら回答する

この問題を1つ目の解法で解こうとすると、毎回DFSをして愚直に回答する方法が思い浮かぶ。
しかし、この方法だとO(N²)になり、時間がかかりすぎてしまう。

2つ目の解法で解くことを考えてみると、頂点ごとにどんな子要素があるかを表す配列を作る方法が思いついた。
しかしメモリを大量に消費する恐れがあるため、解法としては不十分。

最後に3つ目の解法で解けないか考えてみる。

DFS＋クエリ先読みで解く

木に関する問題はDFSで解けることが多い。 根からDFSを始めて、頂点 $ U_i $まで到達した時に頂点 $ U_i $に関するクエリに回答することにする。

また、距離に関する問題であるため、根からの距離がiであるような頂点がいくつあるか、という配列distを作成しておく。
頂点 $ U_i $に到達したときに、その時点でのdistがどうなっているかを記録しておき、頂点 $ U_i $からその親要素に戻る際にもう一度distを確認して、記録しておいた時との差分を取れば、子要素に距離iの要素がいくつあるかを知ることができる。
ただし、頂点ごとにdistをコピーして保存するとMLEになってしまうため、その頂点のクエリに関係する値だけを記録しておく。

具体例

$ U_i = j, D_i = 2 $というクエリに回答することを考える。
頂点jに到達した時点でdist[2]が3であり、頂点jから親要素に戻るときにdist[2]が5だった場合、頂点jの子要素に根からの距離が2である要素が5 - 3 = 2個あることがわかり、これがクエリの答えになる。

実装

クエリ先読み周りの実装が汚くなってしまった。

void dfs(int n, int p, int d, vector<int> &dist, vector<map<int, pair<int, int>>> &query, vector<vector<ll>> &G) {
    for (auto &q : query[n]) {
        q.second.second -= dist[q.second.first];
    }

    dist[d]++;
    for (auto a : G[n]) {
        if (a != p) { // 親への逆流防止
           dfs(a, n, d+1, dist, query, G);
       }
    }

    for (auto &q : query[n]) {
        q.second.second += dist[q.second.first];
    }
}

void solve(long long N, std::vector<long long> P, long long Q, std::vector<long long> U, std::vector<long long> D) {
    // グラフ構築
    vector<vector<ll>> G(N);
    REP (i, N-1) {
        G[P[i]-1].push_back(i+1);
    }

    // 根からの距離
    vector<int> dist(N, 0);
    // クエリを保管する配列 map.key: index, map.value.first: D[i], map.value.second: 答え
    vector<map<int, pair<int, int>>> query(N);
    REP (i, Q) {
        query[U[i]-1].insert({i, {D[i], 0}});
    }
    
    dfs(0, -1, 0, dist, query, G);

    // クエリの順番を元に戻す
    vector<int> ans(Q, 0);
    REP (i, N) {
        for (auto q : query[i]) {
            ans[q.first] = q.second.second;
        }
    }

    // 出力
    REP (i, Q) c(ans[i]);
}

感想

公式解説の方法も勉強しておきます。