ARC228 B - Minimize Topological Order

最小トポロジカル

考え方

まず、$f(T)$ の作り方を考える。
これは、以下のように最良優先探索で貪欲に処理すればよい。

最初に順列は空にしておく
その頂点は未採用だが親は採用済、という頂点一覧を作る
- 最初の $1$ 回は、根のみとする
その中で、最も番号が小さいものを採用し、順列の末尾に追加する
未採用頂点が残っていたら、頂点一覧の作成に戻る

ということは、それが指定の順列 $P$ に一致する根付き木 $T$ は、以下のように作られる。

最初は、順列の先頭を根とする $1$ 頂点のグラフとする
2番目以降について、順に頂点を以下の手順で追加する
- 追加しようとしている頂点より大きい番号の頂点を最後に追加したタイミングを求める
  - 今まで追加したものが全て小さい番号だった場合、根をおいたタイミングとする
- そのタイミング以後（そのものでもよい）に追加した頂点を $1$ つ選ぶ
- それを親とし、追加する頂点をその子とする

よって、この作り方でコストが最小になるような親の選び方を求める問題ということになる。

コストは、$i$ 番の頂点に $k$ 個の子がいると、$A_i\times k^2$ だけかかる。
つまり、既に $k$ 個の子を持っているところに $1$ つ追加するコストは $A_i\times (2k+1)$ である。
この毎回のコストが最小であるように選ぶ貪欲法が成立するか考えよう。

頂点 $P_i$ を追加するときに、その場でコスト最小である親が $X$ だったとしよう。
そこであえて $Y$ を親に選択してみて、これで得することがあるかどうかを考える。

まず、それ以後誰も $X$ を親に選ばなかった場合。
その場合、単純に $P_i$ の親を $Y$ から $X$ に変更しただけの解が、同等以下のコストになる。

一方、それ以降に誰かが親に $X$ を選んだ場合。
次に $X$ を親に選んだ頂点を $P_j$ とする。
この状況では、以下が成立している。

$X$ 追加以後、$P_i$ 追加まで、$P_i$ より大きい数は追加されていない
$Y$ 追加以後、$P_i$ 追加まで、$P_i$ より大きい数は追加されていない
$X$ 追加以後、$P_j$ 追加まで、$P_j$ より大きい数は追加されていない
追加順を考慮すると、これらから以下も成立することがわかる。
$Y$ 追加以後、$P_j$ 追加まで、$P_j$ より大きい数は追加されていない
つまり、実は $P_j$ 追加時に、親の選択肢に $Y$ も存在している。
したがって、$P_i$ の親を $X$ に、$P_j$ の親を $Y$ に、交換してもコストは変わらない。

以上より、あえてその場で損をすることで後で得するパターンは存在しない。
よって、毎回の追加コストが最小であるように選ぶ貪欲法でよいことがわかった。
したがって、$O(N^2)$ でよければこれで解けた。

あとは、高速化である。

まず、選択肢の範囲探索の高速化。
これは、以下のようにすればよい。
空の stack を用意し、順列を後ろから見て、順に以下を実行する。

stack の中に、次に処理する値より小さいものがあれば、すべて取り除く
- 取り除いたタイミングが、取り除いた数の方の選択肢範囲の先頭
stack に処理する値を追加する
- 位置も一緒に入れると、記録がとりやすい
  最後まで stack の中に残りっぱなしだったものは、先頭から選べる。
  末尾は、調べるまでもなく自身の $1$ つ前である。
  これで $O(N)$ ですべての選択肢範囲を調べられる。

そして、コストが最も安い親を選ぶ方法。
これは最小値を求めるsegment木でよい。
頂点番号順でなく順列にある順に、「次にそれを親に選ぶとコストはいくつか？」を入れておく。
選択肢の範囲と同じ範囲で、最小値がある位置が、そのときに親にするべき頂点である。
コスト情報更新も込みで、$O(N\log N)$ で処理できる。

よって、全体として $O(N\log N)$ まで高速化でき、間に合う。

入力例2での動作

入力を受け取る。

n: 5
p: {3, 1, 5, 2, 4}
a: {4, 1, 1, 2, 1}

以下、順列上の位置は 0-indexed で表し、頂点番号は入力のまま表す。

まず、各位置について親として選べる範囲の左端を求める。
順列 $P=\{3,1,5,2,4\}$ を後ろから stack で処理すると、次のようになる。
stack は左から順に底から上を表す。

i=4: 4 を追加
stack: {4}

i=3: 2 を追加
stack: {4, 2}

i=2: 2, 4 を取り除き、5 を追加
b[3] = 2
b[4] = 2
stack: {5}

i=1: 1 を追加
stack: {5, 1}

i=0: 1 を取り除き、3 を追加
b[1] = 0
stack: {5, 3}

最後まで残った頂点については先頭から選べるので、b[2]=0 となる。
b[0] は使用しないダミーとして $-1$ とする。
各位置の左端は次のようになる。

b: {-1, 0, 0, 2, 2}

次に、segment木の位置 $i$ に、頂点 $P_i$ を次に親として選ぶ場合の追加コストを入れる。
最初はどの頂点も子を持たないので、各位置の値は $A_{P_i}$ そのものである。

segment木: {1, 4, 1, 1, 2}
累積コスト: 0

$i=1$、頂点 $1$ を追加する。
親候補は位置区間 $[b[1],1)=[0,1)$ にある頂点 $3$ のみ。
追加コスト $1$ で頂点 $3$ を親にする。

頂点 $3$ の子は $1$ 個になったので、次に子を追加するコストは $1\times(2\times1+1)=3$ となる。

segment木: {3, 4, 1, 1, 2}
累積コスト: 1

$i=2$、頂点 $5$ を追加する。
親候補は位置区間 $[0,2)$ にある頂点 $3,1$ で、追加コストはそれぞれ $3,4$。
よって頂点 $3$ を親にする。

頂点 $3$ の子は $2$ 個になったので、次の追加コストは $1\times(2\times2+1)=5$ となる。

segment木: {5, 4, 1, 1, 2}
累積コスト: 4

$i=3$、頂点 $2$ を追加する。
親候補は位置区間 $[2,3)$ にある頂点 $5$ のみで、追加コストは $1$。
頂点 $5$ を親にする。

segment木: {5, 4, 3, 1, 2}
累積コスト: 5

$i=4$、頂点 $4$ を追加する。
親候補は位置区間 $[2,4)$ にある頂点 $5,2$ で、追加コストはそれぞれ $3,1$。
よって頂点 $2$ を親にする。

累積コストは $6$ となるので、答えは $6$ である。

注意点

コストや答えは int 型からはみ出る。
long long 型を用いること。

別解

特になし。