エクサウィザーズ 2019 E Black or White

atcoder python プログラミング

エクサウィザーズ 2019 E Black or White の解説。

コンテスト中に解けたのが嬉しかったので、ついでに解説を書く。

問題はこれ
E - Black or White。
解説はこれ
https://img.atcoder.jp/exawizards2019/editorial.pdf。

解説

まず、 i 回目までに B 個の黒が全て出る確率 b_i を求める。
これは、 k ( \leq i) 回目で最後の黒が出る確率の和になる。
k回目に、最後の黒が出る場合の数は、 k-1 回目までに黒がちょうど B-1 回出る数に等しい。
つまり、 \binom{k-1}{B-1} になる。
また、 k ( < B) 回目までで黒が全て出ることはない。
まとめると、
b_i = \begin{cases}\sum_{k = B}^{i} \frac{\binom{k-1}{B-1}}{2^k} & (i \geq B)\\ 0 &(i < B)\end{cases}
となります。

同様に、 i 回目までに W 個の白が全て出る確率 w_i は、
w_i = \begin{cases}\sum_{k = W}^{i} \frac{\binom{k-1}{W-1}}{2^k} & (i \geq W)\\ 0 &(i < W) \end{cases}
で与えられます。

では、i 回目に黒が出る確率を求めます。
これは、i 回目の操作を行う前に黒と白がそれぞれ少なくとも１つあれば、
(i-1\mbox{回目の操作の後、白と黒が少なくとも１つある確率})\times \frac{1}{2}
になります。
(i-1回目の操作の後、白と黒が少なくとも１つある確率)
は、i-1 回目の操作の後、黒または、白が出尽くしているわけではないと言い換えられるので、
1 - b_{i-1} - w_{i-1}
となります。
白または、黒しか無い場合は黒しか見る必要がないので、
w_{i-1}
となります。これは、i-1回目までに白が全部出ていれば、残りは黒しか出ないためです。
というわけで、i 回目に黒が出る確率は、
(1 - b_{i-1} - w_{i-1})\times \frac{1}{2} + w_{i-1}
となります。
後は、頑張って実装します。

解答

Q = 10**9+7
def getInv(N):#Qはmod
    inv = [0] * (N + 1)
    inv[0] = 1
    inv[1] = 1
    for i in range(2, N + 1):
        inv[i] = (-(Q // i) * inv[Q%i]) % Q
    return inv

B, W = map( int, input().split())
Z = getInv(B+W)
I = [1]*(B+W+1)
for i in range(1,B+W+1):
    I[i] = (I[i-1]*Z[i])%Q

F = [1]*(B+W+1)
for i in range(1,B+W+1):
    F[i] = F[i-1]*i%Q
w = 0
b = 0
for i in range(1,B+W+1):
    if i - 1 < B and i - 1 < W:
        print(I[2])
    elif i - 1 >= W and i - 1< B: #i > B
        w += F[i-2]*I[i-1-W]%Q*I[W-1]%Q*pow(I[2],i-1,Q)%Q
        w %= Q
        print(((1 -  w)%Q*I[2]%Q + w)%Q)
    elif i - 1 >= B and i - 1 < W: #i > W
        b += F[i-2]*I[i-1-B]%Q*I[B-1]%Q*pow(I[2],i-1,Q)%Q
        b %= Q
        print((1-b)%Q*I[2]%Q)
    else:
        w += F[i-2]*I[i-1-W]%Q*I[W-1]%Q*pow(I[2],i-1,Q)%Q
        w %= Q
        b += F[i-2]*I[i-1-B]%Q*I[B-1]%Q*pow(I[2],i-1,Q)%Q
        b %= Q
        print(((1 -  w - b)%Q*I[2]%Q + w)%Q)
おまけ

mod の逆元を求める記事書いたと思ったけど、書いてなかった。