Zenn
💭

ABC239 E - Subtree K-th Max C++解答例

2022/02/24に公開
C++
#
競技プログラミング
AtCoder
#
競プロ
tech

AtCoder Beginner Contest 239 E - Subtree K-th Max を C++で解きます。

問題

問題文をAtCoder のページより引用します。

問題文

N頂点の根付き木があります。頂点には1からNの番号がついており、根は頂点1です。
i番目の辺は頂点A_iとB_iを結んでいます。
頂点iには整数X_iが書かれています。

Q個のクエリが与えられます。i番目のクエリでは整数の組(V_i, K_i)が与えられるので、次の問題に答えてください。

  • 問題：頂点V_iの部分木に含まれる頂点に書かれた整数のうち、大きい方からK_i番目の値を求めよ。

制約

  • 2 \leq N \leq 10^{5}
  • 0 \leq X_i \leq 10^{9}
  • 1 \leq A_i, B_i \leq N
  • 1 \leq Q \leq 10^{5}
  • 1 \leq V_i \leq N
  • 1 \leq K_i \leq 20
  • 与えられるグラフは木である
  • 頂点V_iの部分木は頂点をK_i個以上持つ
  • 入力に含まれる値は全て整数である

解答例

クエリの解を事前に計算して保持しておく

任意の頂点Vの部分木に含まれる頂点に書かれた整数を降順にソートした配列P_Vを事前に計算しておけば、各クエリについて\mathcal{O}(1)で答えることができます。
しかし、この方法では、全てのノードが一列に並んだような根付き木に対して最悪計算量が\mathcal{O}(N^2)になってしまうため、時間内に解くことができません。

制約を見ると、各クエリにおけるK_iの最大値は20であることがわかります。
これは、部分木についての配列P_Vに対して、先頭から高々20までの要素しか参照しないことを表しています。
したがって、全てのクエリのK_iの最大値をKとおいたとき、任意の頂点Vの部分木に含まれる頂点に書かれた整数のうち、大きい順にK個入った配列をP_Vとして記録すればよいことになります。

葉から順に探索を行う

部分木に含まれる整数の集まりを求めるので、根から求めるよりも葉から求めていった方が効率的です。
配列P_Vを葉から順に求めるとします。m個の子を持つある頂点Vについて、その子らC_1, C_2, \ldots, C_mについてP_{C_j} (1 \leq j \leq m)が求まっていると仮定します。
このとき、全ての配列P_{C_j}をマージして配列Cを用意します。すると、これを降順にソートして、先頭から最大でK個までの要素を取ることで、配列P_Vを作ることができます。

自分の子らについての値を全て求めてから自分についての値を求める、という操作は、深さ優先探索の帰りがけ順で計算を行うことで実現できます。

子についての配列P_{C_j}の要素数は高々20であるので、ある頂点Vについての計算量はmK\log{mK}となります。
各頂点が持つ子の数の総和はNを上回ることは無いので、全ての頂点についてP_Vを求めるときの計算量は\mathcal{O}(NK\log{NK})となります。

プログラム実装例

C++のプログラム実装例を以下に示します。

e.cpp
#include <iostream>
#include <vector>
#include <algorithm>

// 深さ優先探索
void dfs(std::vector<std::vector<int64_t>>& graph, std::vector<int64_t>& X,
         std::vector<std::vector<int64_t>>& P, int64_t K, int64_t current_node,
         int64_t previous_node) {
  // 帰りがけ順に探索する
  // まずは自分の子らについて探索する
  for (int64_t next_node : graph.at(current_node)) {
    // 根に向かう方向に逆流しないようにする
    if (next_node == previous_node) continue;

    // 子らについて探索を行う
    dfs(graph, X, P, K, next_node, current_node);
  }

  // 帰りがけ順の処理。この時点で自分の部分木についての探索は終了している
  // 子らの配列をマージした結果を格納するための配列
  std::vector<int64_t> cache;
  // 自分が持っている整数も含める
  cache.push_back(X.at(current_node));
  // 子らの配列をcacheにコピーする
  for (int64_t next_node : graph.at(current_node)) {
    std::copy(P.at(next_node).begin(), P.at(next_node).end(),
              std::back_inserter(cache));
  }

  // 降順にソートする
  std::sort(cache.begin(), cache.end(), std::greater<int64_t>());

  // 大きい方からK個までを記録する
  int64_t size = cache.size();
  for (int64_t i = 0; i < std::min(K, size); i++) {
    P.at(current_node).push_back(cache.at(i));
  }
}

int main() {
  // 入力
  int64_t N, Q;
  std::cin >> N >> Q;

  std::vector<int64_t> X(N);
  for (int64_t i = 0; i < N; i++) {
    std::cin >> X.at(i);
  }

  std::vector<std::vector<int64_t>> graph(N);
  for (int64_t i = 0; i < N - 1; i++) {
    int64_t a, b;
    std::cin >> a >> b;
    // 0-indexedにする
    graph.at(a - 1).push_back(b - 1);
    graph.at(b - 1).push_back(a - 1);
  }

  std::vector<int64_t> v(Q);
  std::vector<int64_t> k(Q);
  for (int64_t i = 0; i < Q; i++) {
    std::cin >> v.at(i) >> k.at(i);
    // 0-indexedにする
    --v.at(i);
  }

  // 各クエリのKiの中の最大値
  int64_t K = *std::max_element(k.begin(), k.end());
  // クエリの答えを格納する配列
  std::vector<std::vector<int64_t>> P(N);
  // 深さ優先探索開始
  dfs(graph, X, P, K, 0, -1);

  // 各クエリに対する答えの出力
  for (int64_t i = 0; i < Q; i++) {
    std::cout << P.at(v.at(i)).at(k.at(i) - 1) << std::endl;
  }

  return 0;
}

実際に提出したコードはこちら。

GitHubで編集を提案

Discussion