もうハマらない確率・期待値DP

競プロで確率DPや期待値DPの問題で、いつも以下の点でハマっているので、整理して記事化しておくことにしました。

• DPの遷移で混乱する
• ある確率で成功する事象について、成功するまでの試行回数の期待値を毎回1から計算して時間がかかる

想定している問題は以下のような問題です。

本記事の前提

• 本記事では、確率や期待値をDPで解く問題を扱います
  • 確率DPの変数をprob、期待値DPの変数をexで表します
• 期待値は、状態遷移時に発生する何らかのコストやスコアのようなものを想定します(例：サイコロを振った回数、かかった費用)
• 状態iから状態jに遷移について、
  • 遷移確率をp[i][j]で表します(各iについて、p[i][j]をj全体で足した総和は1)
  • その時にかかるコストをcost[i][j]と表します
• 本記事に出てくるコード例は擬似コードです

DPの遷移

典型的によく使うのは以下のパターンになります。期待値については、ゴールから逆順にもらった方がハマりにくいと思われます。

• 確率DP
  • 「配る」DP
• 期待値DP
  • ゴールから逆順に「もらう」DP
  • 「配る」DP

それぞれのパターンについて整理します。

確率DP (「配る」DP)

状態iにいる確率prob[i]がわかっている時に、後続の状態jの確率prob[j]は以下のように更新できます。

for (i = start; i < goal; i++) {
  for (j : iの遷移先) {
    prob[j] += prob[i] * p[i][j]
  }
}

開始地点を0とすると、初期状態はprob[0] = 1.0と初期化します。

期待値DP (ゴールから逆順に「もらう」DP)

状態iからゴールに到達するまでにかかるコストの期待値ex[i]は、それより先の状態j(j>i)からゴールに到達するまでにかかるコストの期待値ex[j]がわかっていれば、以下のように更新できます。

for (i = goal; i >= start; i--) {
  for (j : iの遷移先) {
    ex[i] += (ex[j] + cost[i][j]) * p[i][j]
  }
}

初期値は ex[goal] = 0.0で初期化すれば良いです。

この遷移でスタート地点まで更新すると、ex[start]がスタート地点からゴール地点までにかかるコストの期待値ということになります。

これの意味するところを図解してみます。

まず、そもそもの期待値の求め方を整理しておきましょう。

期待値は確率とコストの積を足していくと得られます。つまり、確率を横軸に、コストを縦軸に取ると、以下の図のように面積の総和が期待値ということになります。

オレンジの部分の面積1.8と、青の部分の面積1.6を足して、期待値3.4が求められます。

これを念頭に置いて期待値DPを考えてみます。

以下のような状態遷移があったとしましょう。

この時、各地点からゴールに到達するまでにかかるコストの期待値を求めることにします。

まずゴール地点は、すでにゴールに到達していてこれ以上コストがかからないので期待値は0です。つまり棒グラフは空です。

次に状態Cを考えると、Cからはゴールに進む1択で、その時にかかるコストは2です。つまり状態Cからゴールに到達するまでの期待値は、1.0 * 2 = 2.0です。棒グラフは横幅一杯(確率1.0)、高さ=2(コスト=2)の箱で表現され、面積が2.0となります。

状態D、状態Eについても同様に求まります。

次に状態Aを考えます。この状態からは状態Cと状態Dにそれぞれ確率0.6、0.4で進み、それぞれコストが1、3かかります。状態Cに進んだ場合かかるコストはA→Cのコスト1に加えて、状態Cからゴールまでのコストの期待値2.0の合計3.0がかかることになります。これが確率0.6で発生するので、この部分の期待値は3.0 * 0.6 = 1.8です。(図の薄いオレンジと濃いオレンジの箱)

状態Cから状態Dに進んだ場合も同様に、(3.0 + 1.0) * 0.4 = 1.6の期待値になります(図の薄い青と濃い青の箱)

つまり、状態Cからゴールまでにかかる期待値は、1.8 + 1.6 = 3.4となります。これはAのグラフの面積の全体になります。

状態Bからの遷移先は3つありますが、これも同様に求めることができます。

期待値DPのアルゴリズムは、このように遷移先の期待値と遷移の確率・遷移時のコストを使って更新していたわけです。

期待値DP (「配る」DP)

先ほどのゴールから逆順にもらうDPでは、期待値のみのDPで済んでいたのに対して、配るDPでは確率も一緒にDPしていくことになるので、もらうDPが可能であればそちらでやるのがオススメです。

しかし、配るDPが必要になった場合についても説明します。

以下のような状態遷移があったとします。

今回は、最終的にゴールした際の期待値に対する各状態での寄与分をDPしていくことになります。これは正確には期待値というわけではないので、変数名はcontribで表します。

ややこしいので、まずは図解から始めます。

左側の図が状態A,Bまで進めた時の状況、右側の図が状態C,D,Eまで進めた時の状況を表します。

これをゴールまで進めていくと、横幅はどんどん細かくなりますが、最終的な面積が期待値となります。

このイメージを念頭に、各状態における期待値貢献分の遷移を整理してみます。

for (i = start; i < goal; i++) {
  for (j : iの遷移先) {
    prob[j] += prob[i] * p[i][j]
    contrib[j] += (contrib[i] + cost[i][j] * prob[i]) * p[i][j]
  }
}

先ほどの図式で考えると、probは横幅、costは高さ、contribは面積なので、この式は以下のことをしていることになります。

• iでの面積contrib[i]に対して、
• 高さcost[i][j](iからjへの遷移時のコスト)、幅prob[i](ex[i]の幅)の長方形を積み重ねて、
• iからjへの遷移確率p[i][j]を掛けて横幅を調整し、状態jでの面積を求めた

初期状態は、prob[0] = 1.0、contrib[0] = 0で初期化すれば良いです。

ここでちょっと気になるかもしれないのは、棒グラフの図でDの状態では、A→Dと進んだパターンと、B→Dと進んだパターンで高さの違う長方形になっていますが、今のやり方だと面積だけ足してしまっているので、内訳の情報は失われてしまっていることです。
しかし、この後Dからの遷移があった場合、A→Dパターン、B→DパターンのそれぞれにDからの遷移のコストを積み重ねることになりますが、これを別々に管理して積み重ねても、平均化してから積み重ねても、結果の面積は変わらないので、特に問題はありません。

最後に変数と意味を整理しておきます

変数	意味	棒グラフ上の意味
prob[i]	状態iにいる確率	棒グラフの横幅
ex[i]	状態iからゴールまでの期待値	棒グラフの面積(状態iでの横幅を1とする)
contrib[i]	最終期待値に対する状態iの寄与	棒グラフの面積(状態iの横幅をprob[i]とする)
cost[i][j]	状態iから状態jへの遷移のコスト	棒グラフの高さ
p[i][j]	状態iから状態jへの遷移確率	棒グラフの横幅(の比)

成功までの試行回数の期待値

確率pで成功する試行を、成功するまで繰り返した時の試行回数の期待値は
$$\frac1p$$
で即求まります。

以下は導出過程を書いておきますが、結果だけ分かれば良い方は読み飛ばしていただいて構いません。
これは高校で習うよくある計算で求まります。求める期待値を$E$とすると、1回目で成功するケース、2回目で成功するケース、...と順番に足していくことで、
$$E=p+2p(1-p)+3p(1-p)^2+4p(1-p)^3+\dots$$
と書けます。これに両辺$1-p$を掛けると
$$(1-p)E=p(1-p)+2p(1-p)^2+3p(1-p)^3+4p(1-p)^4+\dots$$
となり、最初の式からこの式を引くと、右辺は隣の項同士で引き算ができて全て係数が1となり、
$$pE=p+p(1-p)+p(1-p)^2+p(1-p)^3+\dots$$
となります。この右辺は、

$$
\begin{align}
右辺&=p\{1+(1-p)+(1-p)^2+(1-p)^3+\dots\}\\
&=p\frac{1}{1-(1-p)}\\
&=1
\end{align}$$
となり、最終的に
$$E=\frac1p$$
が得られます。

(おまけ)コンプガチャの試行回数の期待値

上記の結果を使うと、いわゆる「コンプガチャ」の試行回数の期待値も求まります。問題を正確に表現すると、$n$種類の賞品が等確率で当たるくじ引きで、全賞品を入手するまでにかかる試行回数の期待値です。

まず、1回目の試行で必ずどれかの賞品は入手できるので、1種類目を入手するまでの期待値は1です。

次に、その後に2種類目を入手するまでの期待値を求めてみましょう。これは1種類目と同じものを引かなければ良いので、成功確率は$(n-1)/n$です。これの逆数が期待値でしたから、期待値は$n/(n-1)$です。

この調子で、3種類目の期待値は$n/(n-2)$、4種類目は$n/(n-3)$となり、最後の1種類は$n/1$が期待値となります。

これを全て足すと、求める期待値は、

$$E=n\Big(1+\frac12+\frac13+\dots+\frac1n\Big)$$
となります。カッコ内に調和数列が出てきています。