この問題は俗に 01 on Tree と呼ばれるテクニックを利用すると解ける問題です。

まず、宝探しの戦略を \((0, 1, \dots, N)\) の順列 \(q = (q_0, q_1, \dots, q_N)\) を用いて表してみましょう。ここで \(q\) は次のような戦略を意味します。

• \(q_0 = 0\) である。
• まず、\(q_1\) を探索する。宝が発見されたら操作を終了する。
• 次に、\(q_2\) を探索する。宝が発見されたら操作を終了する。
• \(\vdots\)
• 最後に、\(q_N\) を探索する。

\(q\) が戦略として合法である必要十分条件は次の通りです。

• \(1 \leq i \leq N\) について、\(q\) 内において \(p_i\) が \(i\) よりも左に登場する。

また、便宜上 \(a_0 = 0, S = \sum_{i=1}^N a_i\) と置くと、戦略 \(q\) を取った時の操作回数の期待値 \(E(q)\) は次式で表されます。

\[ \begin{aligned} E(q) &= \sum_{i=1}^N i \times \frac{a_{q_i}}{S} \\ &= \frac{1}{S} \sum_{i=0}^N i \times a_{q_i} \end{aligned} \]

この問題は \(E(q)\) を最小化する問題です。よって \(\frac{1}{S}\) の部分は定数倍なので一旦無視することができて、\(\sum_{i=0}^N i \times a_{q_i}\) を最小化すればよいことがわかりました。

以上の議論をまとめると、結局次の問題が解ければよいです。

言い換え後の問題 (1)

\(N+1\) 頂点の根付き木が与えられる。 次の条件を満たす順列 \((q_0, q_1, \dots, q_N)\) であって \(\sum_{i=0}^N i \times a_{q_i}\) を最小化するものを求めよ。

• \(i=1, 2, \dots, N\) について、\(i\) の親頂点の番号が \(i\) よりも\(q\) において左に登場する。

この問題は、01 on Tree (リンク) とほとんど同じ問題です。01 on Tree の概要を簡単に書くと次のようになります。

01 on Tree の概要

\(N\) 頂点の根付き木および \(01\) 列 \((V_1,V_2,\dots,V_N)\) が与えられる。
次の条件を満たす順列 \((q_1, \dots, q_N)\) であって \((V_{q_1}, V_{q_2},\dots,V_{q_N})\) の転倒数を最小化するものを求めよ。

• \(i=2, \dots, N\) について、\(i\) の親頂点の番号が \(i\) よりも\(q\) において左に登場する。

さて、説明を簡単にするために言い換え後の問題を 01 on Tree に帰着させましょう。言い換え後の問題 (1) を転倒数を用いて言い換えると次のようになり、この問題は明らかに 01 on Tree の上位互換です。

言い換え後の問題 (2)

\(N+1\) 頂点の根付き木が与えられる。また、\(V_i = (0, 0, \dots, 0, 1)\) (ここで \(0\) は \(a_i\) 個並んでいる) が頂点 \(i\) に書かれている。
次の条件を満たす順列 \((q_0, q_1, \dots, q_N)\) であって、\(V_{q_0}, V_{q_1}, \dots, V_{q_N}\) をこの順に結合してできる列の転倒数を最小化するものを求めよ。

• \(i=1, 2, \dots, N\) について、\(i\) の親頂点の番号が \(i\) よりも\(q\) において左に登場する。

この問題は次の手順で解くことが出来ます。(本解法は 01 on Tree の editorial を参考にして書きました)
まず、\(C_{0,i}\) を頂点 \(i\) に書かれている数列に含まれる \(0\) の個数、\(C_{1,i}\) を頂点 \(i\) に書かれている数列に含まれる \(1\) の個数とします。そして、\(\frac{C_{0,i}}{C_{1,i}}\) が最大の頂点 (複数ある場合はどれでも良い) を \(v\) とおきます。そして \(v\) の親頂点を \(p\) とおきます。
この時、\(q\) において \(p\) の右隣に \(v\) が来るような並び替えが存在することが確認できます。(そうでない場合 \(v\) を左に動かしても損をしないため) つまり、\(q\) において \(p\) の直後に \(v\) が来ると考えてよいです。そこで、頂点の親子関係を保ちながら頂点 \(p\) と頂点 \(v\) を縮約した新たな木を作り、その木において問題を解いても良いです。(ここで縮約した頂点には \(V_p\) と \(V_v\) をこの順に結合した数列を書きこみます。) この操作を頂点数が \(1\) 頂点になるまで行えば問題を解くことが出来ます。
以上の方法によって正答を得るアルゴリズムを得ることが出来ましたが、\(v\) を探して \(p\) とマージする操作を愚直に行うと \(\mathrm{O}(N^2)\) 掛かってしまいます。この部分については priority queue と Union-Find を利用すれば計算量を \(\mathrm{O}(N \log N)\) に落とすことが出来ます。(この部分は容易ではないですが青～黄 diff 程度なのでここでは省略します) また、転倒数の計算については頂点をマージする際に発生する分を適宜管理することで高速に計算できます。

以上の内容を適切に実装することでこの問題を \(\mathrm{O}(N \log N)\) 程度で解くことが出来て、これは十分高速です。