解法メモ AtCoder Grand Contest 044 E - Random Pawn

問題文：[AtCoder Grand Contest 044 E - Random Pawn](https://atcoder.jp/contests/agc044/tasks/agc044_e)

# 準備

まず、$A_i$が最大である位置$i$にポーンがいる時は、ゲームを終えるのが最適。$A_i$で円周を切り離して直線上に並べ、$A_i$が両端になるようにし、$0$から$N$までの番号を左端から振り直す。つまり、$A_0=A_N＝元A_i$ となるようにする。すると、両端に到達した時はゲームを終えるのが最適になり、ポーンが両端からはみ出る行動を考える必要がなくなる。

位置$i$から開始した時の期待利得を$E_i$とする。$E_0=E_N=A_0$である。$E_i$が全て求まれば、$i=0,...,N-1$の範囲での$E_i$の平均が答になる。

# 期待利得の関係式

現在$i$にいるとする。その状態から得られる期待利得$E_i$は、

- ゲームを終えるなら $A_i$
- ゲームを続けるなら $-B_i+(E_{i-1}+E_{i+1})/2$

なので、

```math
E_i = max \left\{ A_i, -B_i+(E_{i-1}+E_{i+1})/2 \right\}
```

となる。
<details><summary>失敗談</summary>
配列$E$を適当に初期化してから、  
$E_{new}[i] = max \left\\{ A_i, -B_i+(E_{old}[i-1]+E_{old}[i+1])/2 \right\\}$  
を繰り返し計算すると、意外とあっという間に収束してくれるのでは…という作戦をはじめ考えたが、$B$が全て$0$の時を想像してみると、端の値が距離10万離れた中央にまで誤差$10^{-10}$で伝搬しなければならず、無理そう。
</details>

# ひもの高さ計算法

遠く離れた2つの位置を$l,r$とする。閉区間$[l,r]$の内部ではゲームを続けるのが最適で、その1つ外側（つまり位置$l-1, r+1$）ではゲームを終了するのが最適である場合を考える[^a]。

この時、閉区間$[l,r]$の内部では$E_i$の関係式の$max$が取れ、以下の式になる。

```math
E_i = -B_i+(E_{i-1}+E_{i+1})/2 \ \ \ \ \ \ \cdots 式(1)
```

また、端は値が確定している。

```math
E_{l-1}=A_{l-1}, \ \ E_{r+1}=A_{r+1} \ \ \ \ \ \ \cdots 式(2)
```

ここで、閉区間$[l,r]$での$E_i$の値の様子を想像してみる。移動するたびにペナルティを払うので、両端から遠いほどたくさんのペナルティを払い、$E_i$の値が小さくなる。ちょうど、両端にひもを渡してたるませたような形になる。実際式の上でも、$E_i$は両隣の平均から$B_i$だけ小さいので、どの位置も下に凸であることが分かる。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/605081/5663e949-ceb8-eb50-d3a5-d4dbb9f31d07.png)

さて、このひもの両端の高さから途中の高さを計算できないと問題を解くことができなそうなので、計算できるかを考える。繰り返し計算で収束させる作戦ではだめそうなので、直接求められないか。つまり、式(1)と式(2)を整理して並べてできる、以下のような、変数が$E_l,...,E_r$の$m$個($m=r-l+1$)である連立1次方程式が解けるか。ガウスの消去法は$O(m^3)$なので使えない。

```math
\begin{array}{rrrrrrrrrl}
2E_l&- E_{l+1}&         &         &       &       &         &         &         &= -2B_l+A_{l-1}\\
-E_l&+2E_{l+1}&- E_{l+2}&         &       &       &         &         &         &= -2B_{l+1} \\
    &- E_{l+1}&+2E_{l+2}&- E_{l+3}&       &       &         &         &         &= -2B_{l+2} \\
    &         & -E_{l+2}&+2E_{l+3}&-E_{l+4}&      &         &         &         &= -2B_{l+3} \\
    &         &         &- E_{l+3}&      &\ddots\ &         &         &         &  \vdots \\
    &         &         &         &\ddots\ &      &- E_{r-2}&         &         &  \vdots \\
    &         &         &         &      &-E_{r-3}&+2E_{r-2}&- E_{r-1}&         &= -2B_{r-2} \\
    &         &         &         &       &       &- E_{r-2}&+2E_{r-1}&- E_r    &= -2B_{r-1} \\
    &         &         &         &       &       &         &- E_{r-1}&+2E_r    &= -2B_r+A_{r+1}
\end{array}
```

行列で書くと以下。

```math
\left( \begin{array}{ccccc}
2  & -1 &    &      &      &    &    &    \\
-1 &  2 & -1 &      &      &    &    &    \\
   & -1 &  2 & -1   &      &    &    &    \\
   &    & -1 &      &\ddots&    &    &    \\
   &    &    &\ddots&      & -1 &    &    \\
   &    &    &      &   -1 &  2 & -1 &    \\
   &    &    &      &      & -1 &  2 & -1 \\
   &    &    &      &      &    & -1 &  2 \\
\end{array} \right)
\left(
\begin{array}{c}
E_l \\ E_{l+1} \\ E_{l+2} \\ \vdots \\ \vdots \\ E_{r-2} \\ E_{r-1} \\ E_r
\end{array}
\right) = \left(
\begin{array}{c}
-2B_l+A_{l-1}\\ -2B_{l+1}\\ -2B_{l+2}\\ \vdots\\ \vdots\\ -2B_{r-2}\\ -2B_{r-1}\\ -2B_r+A_{r+1}
\end{array}
\right)
```

これが何と$O(m)$で解けてしまう。まず、$m$本の式を全部足してみたくなる。出てくる式は以下。

```math
E_l+E_r=-2\sum_{i=l}^{r} B_i +A_{l-1}+A_{r+1}
```

また、上から1行目の式に1をかけ、2行目に2をかけ、…、というように、$1$から$m$を上の式から順にかけて足しこむと、$E_r$以外の変数が全て消え、以下の式になる[^b]。
（以下の式の$\sum$は$i=l$から始まるので、かける値は$(i-l+1)$になることに注意）

```math
\begin{align}
(m+1)E_r&=-2\sum_{i=l}^r (i-l+1)B_i +A_{l-1}+mA_{r+1} \\
        &=-2\sum_{i=l}^r iB_i +2(l-1)\sum_{i=l}^r B_i +A_{l-1}+mA_{r+1}
\end{align}
```

これで$E_l$と$E_r$が求まった。$\sum_{i=l}^{r}iB_i$と$\sum_{i=l}^{r}B_i$は、予め累積和を計算しておけば$O(1)$で求まる。後は、$E_{l+1},E_{l+2},...$というように、端から順に値を決めていくことができ、$m$個の変数の値を$O(m)$で計算できる。

しかし、この方法だと、区間の中央付近の値を1つだけ知りたい場合でも、$O(m)$の計算量が必要である。そこで更に、区間内の任意の値1つを$O(1)$で計算する方法を考える。

今、$E_k$ $(l \le k \le r)$を求めたいとする。先ほどの$E_r$を求めた方法を応用すると、うまいこと任意の位置$k$の値が一発で求められる。上から$k'$番目($k'=k-l$)に以下の式があるが、

```math
-E_{k-2} + 2E_{k-1} - E_k = -2B_{k-1}
```

この式を1倍、その上の式を2倍、その上を3倍…として一番上まで足しこむと、$E_k$と$E_l$しか残らないので、$E_k$が$O(1)$で求められる。整理すると以下の式になる。

```math
E_k = - 2\sum_{i=l}^{k-1}iB_i + 2k\sum_{i=l}^{k-1}B_i -k'A_{l-1} + (k'+1)E_l
```

これで、ひもの両端の高さから、任意の位置$k$のひもの高さ$E_k$を$O(1)$で求められるようになった。

# ひもとフック

元の問題に戻る。前述のとおり、位置$i$でゲームを終えるのが最適なら$E_i=A_i$で、ゲームを続けるのが最適なら$E_i=-B_i+(E_{i-1}+E_{i+1})/2$である。これをひもの高さで言い換えると、ゲームを終えるのが最適な位置$i$ではひもの高さを$A_i$に固定し、それ以外の位置ではひもを自然にたるませる、ということになる。以下の図は、赤丸が$(i,A_i)$であり、いくつかの位置$i$で赤丸にひもを固定した様子である。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/605081/8f5f3e65-188a-44f5-6200-4875d200615c.png)

各赤丸にフックがあると考え、固定する ＝ フックにひもを引っかける、と考える。最適なとき、ひもはフックにかかっているか、フックの**上を通過している**はず。逆に、ひもが全てのフックの上を通過（引っかかる場合を含む）するとき、それ以上ひもの位置を上昇させることができないので、最適である。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/605081/ce506571-d7e1-e265-5526-36bc86495261.png)

このときにどのフックに引っかかっているかが分かれば、フックに引っかかっている位置での$E_i$は$A_i$、たるんでいる位置での$E_i$は上述の「ひもの高さ計算法」で計算できるので、答が求まる。

# 作戦1：フックにかけていく作戦

初期状態は両端$i=0$と$i=N$だけ固定されているとし、残りはひもがたるんだ状態であるとする。ひもより高いところにフックがあったら、その場所を上に引き上げてフックにかける。すると、ひも全体としては、変化しない部分と、上昇する部分しかない。これを繰り返して、ひもの高さを最大限に上げた時が答になるはず。

一度フックにかけた場所でも、周囲を上に引き上げた結果、ひもがフックから「外れる」こともあり得る。そこで、外れないことが確定する場所から確定していくとうまくいきそう。おそらく、ひもとフックの高さの差が大きいところから順に決めていくとよいと思われる。

一見、決めるたびに区間が半分になるから$O(N\log N)$で解けるかと思ったが、決まるのが端から順だと$O(N^2)$になってしまう。作戦1ではどうしても計算量を落とすことができなかった。
<details><summary>失敗談</summary>
$A_i$の大きい順に確定していけばよいかと思いきや、WAだった。自分より低いところが決まることで、「外れる」ことがあり得るためだった。
</details>

# 作戦2：フックから外していく作戦

逆に、フックにかけないことが確定したところを捨てていく方向で考えてみる。左端から順に1つずつフックにかけてみる。

初期状態は、$i=0,1$の2か所だけフックにかかっている。$i=2$より右では、ひもが下方に垂れ下がっているとする。その状態から、$i=2,3,...$と順にフックにかけていくとする。新たにフックにかけることで、それまでかかっていたフックが外れたら、二度とそのフックにかかることはない。ひもの位置が下がることはないため。

今、位置$i-1$まで進んだとする。位置$i-1$は必ずフックにかかっている。そこより左は、かかっているところ、かかっていないところがある。どこがかかっているかは把握しているとする。その状態で、位置$i$のフックにひもをかける。ひもをかけるというよりは、位置$i$のひもの高さを徐々に上げていくことを想像する。まず位置$i-1$のフックが外れ、更に上げていくとその一つ左が外れ、というように、上げていくにしたがって、フックが右から左に向かって順に外れていく。逆に言うと、もし位置$i-1$のフックが外れなければ、それより左が外れることはない[^c]。

よって、右から（位置$i-1$から）順に外れるか判定し、外れなかったら終了する、というアルゴリズムでよい。外れるか否かの判定は、「ひもの高さ計算法」を使って$O(1)$で可能。例えば位置$j$が外れるか否か判定するときは、位置$j$の次（1つ左）を位置$k$として、位置$k$と位置$i$を両端とし間をたるませたひもを考え、位置$j$でのひもの高さを計算して、$A_j$と比較すればよい。このアルゴリズムなら、1回の操作ごとに、外れるフックが1つ確定 or 次のフックに進むのいずれかなので、$O(N)$で動作する。

以上のアルゴリズムは、凸包を得るアルゴリズムと流れが同じになっている。公式解説では、配列$C$を使って式変形し、凸包を得る問題そのものに落としてしまうところが鮮やか。いちいちひものことを考えなくてよくなる。

<details><summary>失敗談</summary>
「ひもの高さ計算法」の前半（端から順に計算する方法）で、作戦2を実装すると、TLEだった。$O(N\log N)$になっているかと思いきや、例えば以下のような時に$O(N^2)$になってしまう：
「$i=0$と$i=50000$の$A_i$が大きな値、残りは全部小さな値」
$i=100000$以降、毎回50000の計算量が発生してしまう。
</details>

# 感想

4日間かけ、3度目の挑戦でようやくACできました。自分には、ひもの高さが計算できることに気づくことと、更に任意位置を$O(1)$で計算できるのを発見したことが、最大の興奮ポイントでした。

[^a]: 例えば、$A_{l-1}=A_{r+1}=10^{12}$で、間が全て$0$で、$B$が小さければ、上記行動が最適になりそう。
[^b]: もう一つ式を立てるべくあれこれ足したり引いたり実験し、この重み付き和に行き着いた。
[^c]: 証明はやればできると思うが、ひもを想像することでカバーする。