Python で解く AtCoder M-SOLUTION プロコンオープン C

2019年6月1日に行われた M-SOLUTIONS プロコンオープン C - Best-of-(2n-1) の Python3 での実装例です。 

前提知識

この問題を解くためには、以下の前提知識が必要です。

確率  p で起こる事象が起こるまで試行を続けた際の試行回数の期待値は  p^{-1}

例えば、コインの表が出るまでコインを投げるとしたときの投げる回数の期待値は  2 であり、 1 が出るまでサイコロを振るとしたときのサイコロを振る回数の期待値は  6 です。

証明
確率  p  (0 \lt p \lt 1)で起こる事象が  n 回目で初めて起こるとは、  n-1 回連続で目的の事象以外が起こり  n 回目で目的の事象が起こると言い換えることができます。 そのため、その確率は  p (1-p) ^ {n-1} となります。
よって、期待値の定義より、確率  p で起こる事象が起こるまで試行を続けた際の試行回数の期待値  E は、以下の無限級数となります。

 E = p + 2 \times p(1-p) + 3 \times p(1-p) ^ 2 + \cdots = \Sigma _ {n = 1} ^ {\infty} n \cdot p(1-p) ^ {n-1}

ここで、期待値  E (1-p) をかける

 (1-p)E = \Sigma _ {n=1} ^ {\infty} n \cdot p(1-p) ^ n = \Sigma _ {n=1} ^ {\infty} (n+1) \cdot p(1-p) ^ {n-1}

そして、  E - (1-p)E の式変形よって、期待値  E を公比  (1-p) の無限等比級数とみなせます

 E - (1-p)E = \Sigma _ {n=1} ^ {\infty} n \cdot p(1-p) ^ {n-1} - \Sigma _ {n=1} ^ {\infty} (n+1) \cdot p(1-p) ^ {n-1}
 pE = \Sigma _ {n=1} ^ {\infty} p(1-p) ^ {n-1}
 E = \Sigma _ {n=1} ^ {\infty} (1-p) ^ {n-1}

ここで、  0 \lt 1-p \lt 1 であるため 等比級数の公式 より、

 E = \lim _ {n \to \infty} \frac{1 + (1-p) ^ {n+1}}{1 - (1-p)} = \frac{1}{p}

以上で、確率  p で起こる事象が起こるまで試行を続けた際の試行回数の期待値は  p ^ {-1} であることが証明できました。

方針

期待値の求め方
高橋君がゲームに確率が  \frac{A}{100} 、青木君がゲームに勝つ確率が  \frac{B}{100}、引き分けになる確率が  \frac{C}{100} であるとき、引き分けにならない確率(すなわち、高橋君か青木君のどちらかがゲームに勝つ確率)は  \frac{100 - C}{100} となります。
よって先ほどの前提知識から、高橋君か青木君が1回勝つまでゲームを行った時のゲーム数の期待値は  \frac {100}{100 - C} となり、  M 回の決着がつくまでにゲームを行う期待値は  M \times \frac{100}{100-C} となります。

このゲームの最終的な勝敗は 高橋君か青木君のどちらかが合計で  N 回勝った時に決まります。
ここで、  m 回目の決着がついた時点で高橋君か青木君のどちらかが合計で  N 回勝つときの確率を  X _ m とします。  m 回目の決着がつくまでに行うゲーム数の期待値は  m \frac{100}{100 - C} であるため、最終的な決着がつくまでに行うゲーム数の期待値  E

 E = \Sigma _ {m = N} ^ {2N-1} X _ m \times m \frac{100}{100 - C}

確率  X _ m
 m 回目の決着がついた時点で高橋君か青木君のどちらかが合計で  N 回勝つ確率  X _ m は、引き分けのない条件における  m 回目のゲームで高橋君が合計で  N 回勝つ確率と青木君が合計で  N 回勝つ確率の和と言い換えることができます。
また、引き分けがない条件で高橋君が勝つ確率は  \frac{A}{A+B} であり、同様に青木君が勝つ確率は  \frac{B}{A+B} となります。よって、  X _ m は以下の式で表すことができます。

 X _ m = {} _ {m-1} \mathrm{C} _ {N-1} \left\{\left(\frac{A}{A+B}\right) ^ N \left(\frac{B}{A+B}\right) ^ {m-N} + \left(\frac{B}{A+B}\right) ^ N \left(\frac{A}{A+B}\right) ^ {m-N} \right\}

以上より、求めるべき期待値  E は以下の式で求めることができます。

 E = \frac{100}{100-C} \cdot \Sigma_{m=N} ^ {2N-1} m \cdot {} _ {m-1} \mathrm{C} _ {N-1} \left(\frac{A ^ N B ^ {m-N} + A ^ {m-N} B ^ N}{(A+B) ^ m}\right)

求めるべき整数  R
ただし、出力する答えは期待値  E ではない点に注意してください。

求める期待値は互いに素な整数  P, Q を用いて  \frac{P}{Q} と表せます。
 R \times Q \equiv P \pmod{10 ^ 9+6} となる  0 以上  10 ^ 9+6 以下の整数  R を出力してください。 (この問題の制約下で、このような  R は必ず一意に存在します。)

これは、  E = \frac{P}{Q} であるとき、求めるべき  R は次の式で表されます。
 R = P \cdot Q ^ {-1} \pmod{10 ^ 9+6}
ここで、  Q ^ {-1} Q逆元 です。

逆元は フェルマーの小定理 より、次の式で求めることができます。
 Q ^ {-1} \equiv Q ^ {10 ^ 9 + 6 - 2} \pmod{10 ^ 9 + 6}

すなわち、  R は期待値  E を求める式の分子に 分母の逆元をかけ、最後に  10 ^ 9+6 を法とした剰余計算を行えば求めることができます。(実際には剰余計算は最後にまとめるのではなく、計算量を落とすために項ごとに剰余計算を行います)。

実装