Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 239 E問題 Subtree K-th Max

AtCoder ABC ABC239 E問題 500 pts

問題

提出解答

問題の概要

$N$ 頂点からなる頂点 $1$ を根とする根付き木がある. $j$ 番目の辺は頂点 $A_j, B_j$ を結ぶ. また, 頂点 $i$ には整数 $X_i$ が書かれている.

次の $Q$ 個の問題に答えよ.

頂点 $V_q$ を根とする部分木にある頂点に書かれている整数のうち, $K_q$ 番目に大きい整数は何か?

制約

$2 \leq N \leq 10^5$
$0 \leq X_i \leq 10^9$
$1 \leq A_j, B_j \leq N$
$1 \leq Q \leq 10^5$
$1 \leq V_q \leq N$
$1 \leq K_q \leq 20$
頂点 $V_q$ を根とする部分木の位数 (頂点の数) は $K_q$ 以上

解法

列 $U,V$ に対して, $\operatorname{sort}(U)$ で $U$ を降順に並び替えた列, $U \oplus V$ で $U,V$ を連結させた列とする.

また, 頂点 ${\rm root}$ を根とする根付き木 $T$ に対して, 頂点 $v$ の個全体の集合を $\operatorname{ch}(v)$ とする.

このとき, 各頂点 $v$ に対して, $L_v$ を $v$ を根とする部分木の各頂点に書かれている整数を降順に並べた列とする. すると,

$\displaystyle L_v=\operatorname{sort}\left(\left(\bigoplus_{w \in \operatorname{ch}(v)} L_w \right) \oplus (X_v) \right)$

である.

ここで, $\operatorname{sort}(U \oplus V)=\operatorname{sort}(\operatorname{sort}(U) \oplus V)$ であること及び, $K_{\rm max}=\max(K_1, \dots, K_Q)$ としたとき, $L_v$ の $K_{\rm max}$ 項目より先の項はクエリの解答に影響を与えないことから各 $L_v$ は高々 $K_{\rm max}$ 項だけ持っていれば良い.

このことから計算量 $O(NK_{\rm max} \log K_{\rm max}+Q)$ で求めることができた.