ABC239E - Subtree K-th Max

子ノードの情報を結合してソート

Editorial - Denso Create Programming Contest 2022(AtCoder Beginner Contest 239)の方法．
毎回ソートして間に合う！

考え方

各頂点$v$に対し，「頂点$v$の部分木に含まれる頂点に書かれた整数が降順にソートされたリスト$P_{v}$」を求めておけば，各クエリを$O(1)$で処理できる．

葉から根の方へ親をたどって$P_{v}$を求めていく．

ある頂点$v$を見ているとき，その頂点のそれぞれの子ノード$v_{i}$に対して「その子ノードの部分木に含まれる頂点に書かれた整数が降順にソートされたリスト$P_{v_{i}}$」がわかっているとする．

このとき，$P_{v}$は，子ノードの情報$\{P_{v_{i}}\}_{i}$（と頂点$v$に書かれた整数）をマージして降順にソートすることで求められる．

$1 \leq K_{i} \leq 20$だから，$P_{v}$は大きい方から20個の値を保持しておけば（その親ノードの計算をするのに）十分である．

計算量は，すべての頂点$v$に対し，その子ノード$\{v_{i}\}_{i=1}^{N_{v}}$のリスト（要素数$20$）をマージ（$O(20N_{v})$）してソート（$O(20N_{v} \log(20N_{v}))$）するので，合計で

\begin{aligned}
&O\biggl( \sum_{v} 20N_{v} \log(20N_{v}) \biggr) \\
&= O\biggl( \log \prod_{v} (20N_{v})^{20N_{v}} \biggr) \\
&\leq O\biggl( \log (20N)^{\sum_{v} 20 N_{v}} \biggr) \\
&=20N \log (20N)
\end{aligned}

となり，間に合う．

回答例

頂点を入れるスタックを用意する．はじめ，スタックには根だけが入っているとする．
また，頂点をすでに見たかどうかを管理する（頂点の情報と一緒にスタックに入れる，別なリストを用意る，など）．

スタックからpopした頂点を初めてみる場合は，その頂点を「チェック済み」にした上でスタックにpushする．その後，子ノードを「未チェック」の状態でスタックにpushする．

こうすると，スタックからpopした頂点が「チェック済み」のとき，その部分木は全て「チェック済み」になっている．よって，この処理を使えば，「葉から順に見ていく」ことが可能である．

これは，結局オイラーツアー（Euler tour technique - Wikipedia）の順に見ていることになる．

N, Q = map(int, input().split())
X = list(map(int, input().split()))
G = [[] for _ in range(N)]
for _ in range(N - 1):
  A, B = map(int, input().split())
  A -= 1
  B -= 1
  G[A].append(B)
  G[B].append(A)
  
P = [[x] for x in X]
stack = [(-1, 0)]
seen = [False] * N

while stack:
  par, cur = stack.pop()
  if not seen[cur]:
    seen[cur] = True
    stack.append((par, cur))
    for chi in G[cur]:
      if chi != par:
        stack.append((cur, chi))
  else:
    for chi in G[cur]:
      if chi != par:
        P[cur] += P[chi]
    P[cur] = sorted(P[cur], reverse = True)[:20]
  
for _ in range(Q):
  V, K = map(int, input().split())
  print(P[V - 1][K - 1])

【参考】Submission #29438202 - Denso Create Programming Contest 2022(AtCoder Beginner Contest 239)