ABC 反省文シリーズです。

問題リンク

https://atcoder.jp/contests/abc189/tasks/abc189_f

問題概要

マス $0$ から始めてマス $N$ を目指す双六があります。高橋くんは $1$ から $M$ までの目が出るルーレットを回して進みます。また、この双六のマス $A_i\ (i=1,\ 2,\ ,\dots,\ K)$ は振り出しに戻るマスで、そこに止まるとマス $0$ に戻されます。

高橋くんがマス $N$ より先に着きゴールするまでにルーレットを回す回数の期待値を求めてください。到達不可能ならその旨を報告してください。

制約

• $1\le N,\ M\le 10^5$
• $0\le K\le 10$
• $0\lt A_1\lt \dots\lt A_K\lt N$

考察

まず、$M$ 個以上連続して振り出しに戻るマスが存在したら到達は不可能です。そうでないときを考えます。

見た目からして期待値 DP です。期待値 DP を考えます。期待値を求めるときの定石ですが、初手からその状態になるまでの期待値ではなくその状態からゴール状態になるまでの期待値を考えると往々にして上手くいきます¹。今回も実際に初手からその状態になるまでの期待値を求めると非常に難しくなってしまうので、逆から考えることにします。

さて、$dp_i$ をマス $i$ から $N$ に辿り着くために必要なルーレットを回す回数の期待値と定義します。このとき、「振り出しに戻る」でないマスについては次のような遷移となります。

• $\displaystyle dp_i=1+\frac{1}{M}\sum_{j=1}^M dp_{i+j}$

これは、$1,\ 2,\ ,\dots,\ M$ 個先のマスにそれぞれ $\frac{1}{M}$ の確率で到達し、ルーレットを回す回数はどの場合も $1$ 増えることから分かります。ただし、$dp_i=0\ (i\ge N)$ と便宜上定義します。

そうでないとき、つまり「振り出しに戻る」マスのときはどういう遷移となるでしょうか。「振り出しに戻る」マスに止まるとその次はマス $0$ からやり直しになるため、最初からゴールするまでの期待値そのものが必要なルーレットを回す回数の期待値となります。よって、このときは

• $dp_i=dp_0$

となります。

ここで、期待値を求める際には期待値が満たすべき方程式を解いて求めても良い²ということを思い出します。たとえばサンプル 2 を例にして考えると、

• $1$ が出たら振り出しに戻る。
• $2$ が出たらゴールする。

のそれぞれの起こる確率は $\frac{1}{2}$ です。求めるべき期待値を $E$ と置くと、マス $1,\ 2$ から $2$ まで行くのに必要な回数の期待値はそれぞれ $E,\ 0$ なので、マス $0$ から $2$ まで行くのに必要な回数の期待値 $E$ は $E=\frac{1}{2}(E+0)+1$ を満たすことが分かります。よって、この方程式を解いて $E=2$ とするとこれが答えになっています。

話が逸れましたが、以上のような DP テーブルの遷移をしていくと、$dp_0$ は $dp_0$ の一次式 $dp_0=a\times dp_0+b$ で表されることが分かります。ここで、遷移からも分かりますが $a$ はゴールするまでに振り出しに戻る確率ですので、到達可能であることと $a\lt 1$ となることは同値です。よって、$\displaystyle dp_0=\frac{b}{1-a}$ として期待値を求められます。DP の遷移は $a,\ b$ それぞれを Segment Tree に乗せて行うことができます。

解法

$M$ 個以上連続して振り出しに戻るマスが存在したら到達は不可能です。そうでないときを考えます。

Segment Tree 上で DP を行います。$dp_i(x)=a_i x+b_i$ を、$dp_i(dp_0)$ がマス $i$ からゴールするまでに必要なルーレットを回す回数の期待値となるような一次式と定義します。実装上は $a,\ b$ をそれぞれ別の Segment Tree で保持するか、面倒なら Pair として持つかすれば良いでしょう。
ただし、便宜上 $dp_{N+1}=\dots=dp_{N+M}=0$ と定義します (全部 $0$ なので $N$ より後ろは足切りする実装でも良いです)。

後ろから順に見ていき、マス $i$ が振り出しに戻るマスでないときは

• $\displaystyle dp_i=1+\frac{1}{M}\sum_{j=1}^{M} dp_{i+j}$

振り出しに戻るマスであるときは

• $\displaystyle dp_i=x$

というふうに漸化式が立ちます。答えは $\displaystyle \frac{b_0}{1-a_0}$ です。

実装

		public void Solve()
		{
			var (n, m, k) = sr.ReadValue<int, int, int>();
			var a = sr.ReadIntArray(k);
			var skip = new bool[n + 1];
			for (int i = 0; i < k; i++)
			{
				skip[a[i]] = true;
			}
			var e = new SegmentTree<double>(n + 1, (x, y) => x + y, 0);
			var x = new SegmentTree<double>(n + 1, (x, y) => x + y, 0);
			for (int i = n - 1; i >= 0; i--)
			{
				var ok = false;
				if (m <= 10)
				{
					for (int j = 1; j <= m; j++)
					{
						if (i + j < n + 1 && !skip[i + j]) ok = true;
					}
				}
				else ok = true;
				if (!ok)
				{
					Console.WriteLine(-1);
					return;
				}
				if (skip[i])
				{
					e.Update(i, 0);
					x.Update(i, 1);
				}
				else
				{
					e.Update(i, e.Prod(i + 1, Min(n + 1, i + m + 1)) / m + 1);
					x.Update(i, x.Prod(i + 1, Min(n + 1, i + m + 1)) / m);
				}
			}
			Console.WriteLine(e[0] / (1 - x[0]));
		}

ACコード: https://atcoder.jp/contests/abc189/submissions/19662693

感想

期待値 DP が得意なほうなはずのにこの問題を落としちゃいけないと思いました。なんで後ろから見るっていうのが分からなかったんだろう……ひとつの解法にハマってしまったときのリカバリ方法もちゃんと普段から意識しておかないといけませんね。