日々ｄｒｄｒする人のメモ

今日も明日もｄｒｄｒ

トップ > ProCon > AtCoder: エクサウィザーズ2019 - E問題: Black or White

AtCoder: エクサウィザーズ2019 - E問題: Black or White

ProCon AtCoder

コンテスト中に解けなかった問題。思考力が足りなかった (終)
想定解法メモ。

atcoder.jp

問題概要

{B}個の黒いチョコレートと {W}個の白いチョコレートがある。

黒と白を等確率で選び、選んだ色のチョコレートが存在すれば食べる、をチョコレートがなくなるまで繰り返す。

各 {1 \le i \le B+W} について、 {i}番目に食べるチョコレートが黒である確率をmod {10^9 + 7}で求めよ。

制約

{1 \le B, W \le 10^5}

解法

うまく確率を計算する。


まず、 {i-1}番目までどのような選び方をしても黒と白のチョコレートがそれぞれ1枚以上残る場合、 {i}番目の答えは {\frac{1}{2}}になる。

問題は、片方がなくなりうる場合( {\min(B, W) < i})の確率である。

黒を {b}枚、白を {w}枚食べたことを {(b, w)}と表すとき、確率が {\frac{1}{2}}でなくなるのは、以下の場合である。

  • {(B, k)}の時: 次は必ず白を食べる -> 確率は0
  • {(k, W)}の時: 次は必ず黒を食べる -> 確率は1

よって、 {i}番目で黒を選ぶ確率は、

  • (i番目でどちらも選べる確率) × {\frac{1}{2}} + (i番目で黒しか選べない確率) × 1

が計算できればよいことが分かる。


ここで、 {(B, i-B)}に到達する確率を {p_i}、 {(W, i-W)}に到達する確率を {q_i}とする。

f:id:smijake3:20190331212152p:plain
piとqiのイメージ

これが計算できれば、 {i}番目の解は
(i番目でどちらも選べる確率) × {\frac{1}{2}} + (i番目で黒しか選べない確率) × 1 = {\displaystyle \frac{1-p_i-q_i}{2} + q_i = \frac{1-p_i+q_i}{2}}
と求められる。

この {p_i, q_i}は、以下のように計算できる。
{\displaystyle p_i = \left\{ \begin{array}{ll} 0 & (i < B) \\ p_{i-1} + (\frac{1}{2})^i {}_{i-1}C_{B-1} & (i \ge B) \end{array} \right. }
{\displaystyle q_i = \left\{ \begin{array}{ll} 0 & (i < W) \\ q_{i-1} + (\frac{1}{2})^i {}_{i-1}C_{W-1} & (i \ge W) \end{array} \right. }

階乗とその逆数を前もって計算することで、計算量は {O(B+W)} になる。

実装

提出コード(Python3): Submission #4787508 - ExaWizards 2019

import sys
write = sys.stdout.write
 
B, W = map(int, sys.stdin.readline().split())
MOD = 10**9 + 7
 
# 前計算
L = B+W
fact = [1]*(L+1)
rfact = [1]*(L+1)
for i in range(L):
    fact[i+1] = r = fact[i] * (i+1) % MOD
rfact[L] = pow(fact[L], MOD-2, MOD)
for i in range(L, 0, -1):
    rfact[i-1] = rfact[i] * i % MOD
 
rev2 = pow(2, MOD-2, MOD)


# 各iについて計算していく
p = q = 0
 
L = min(B, W)
r2 = "%d\n" % rev2
for i in range(L):
    write(r2)
base = pow(rev2, L-1, MOD)
 
rB = rfact[B-1]; rW = rfact[W-1]
for i in range(L, B+W):
    base = rev2 * base % MOD
    if i >= B:
        p = (p + (fact[i-1]*rfact[i-B] % MOD) * (rB * base % MOD) % MOD) % MOD
    if i >= W:
        q = (q + (fact[i-1]*rfact[i-W] % MOD) * (rW * base % MOD) % MOD) % MOD
    write("%d\n" % ((1-p+q)*rev2 % MOD))