(C++使用者向けの内容です。)
木DPに帰着するところまでは、公式解説 と同じです。

公式解説と同様、次の要件を満たすデータ構造があればこの問題を解くことができます。

• (多重)集合 $S$ に $X$ を追加する。
• $S$ から $X$ を削除する。
• $S$ の中央値を求める。

これらのクエリは、g++拡張によって使える tree というデータ構造を使うことで効率的に処理することができます。

このデータ構造はC++のsetの上位互換のような感覚で、(多重でない)集合において $k$ 番目の値を返すことや $k$ 未満の数が何個あるかを知ることができます。
参考記事
参考記事

ただし、今回の問題では $A_i$ が同じ値を含む可能性があるのでそのままではうまくいきません。そこで、頂点に書かれた整数と頂点の番号からなるpairを $S$ に追加することで $S$ が多重集合になることはなく、treeを使ってこの問題を解くことができます。
実装例