ABC239 E - Subtree K-th Max C++解答例
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 - 入力に含まれる値は全て整数である
解答例
クエリの解を事前に計算して保持しておく
任意の頂点
しかし、この方法では、全てのノードが一列に並んだような根付き木に対して最悪計算量が
制約を見ると、各クエリにおける
これは、部分木についての配列
したがって、全てのクエリの
葉から順に探索を行う
部分木に含まれる整数の集まりを求めるので、根から求めるよりも葉から求めていった方が効率的です。
配列
このとき、全ての配列
自分の子らについての値を全て求めてから自分についての値を求める、という操作は、深さ優先探索の帰りがけ順で計算を行うことで実現できます。
子についての配列
各頂点が持つ子の数の総和は
プログラム実装例
C++のプログラム実装例を以下に示します。
#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;
}
実際に提出したコードはこちら。
Discussion