ABC189のD問題の解き直し

概要

Atcoder Beginner Contest 189のD問題が、（茶色Diffの問題なのに）手こずってしまったのと、解説だけだとピンと来づらかったので、自分への戒めも兼ねてまとめてみました。

問題

N個の文字列（値は"AND", "OR"のいずれか）が与えられる。
値がTrue or FalseのN+1個の列X（x0, x1, ..., xN）に対して、以下の操作を行った時のXの個数を求める。

• $y0 = x0$
• Siが"AND"のとき、yi = yi-1 && xi
• Siが"OR"のとき、yi = yi-1 || xi ※&&, || はそれぞれ論理積、論理和

公式の解説

公式の解説では、答えをf(S1, ..., SN)とすると、

• SNが"AND"のとき yN-1, xNが一意に定まるため、 $f(S1, ..., SN) = f(S1, ..., SN-1)$
• SNが"OR"のとき xNがTrueならyN-1がなんでも良く、 xNがFalseならyN-1がTrue（一意に定まる）ため、$f(S1, ..., SN) = 2^N + f(S1, ..., SN-1)$

と場合分けができて、Siに応じて、2**(i+1)を加算していくと言った形で答えが出せる。
この方法は、数学の得意な方々には、簡単に導けるのかもしれないが、自分はそうではないので、DPを使って求めます。

DPを使った解法

dpの配列をi番目のyがFalseのときのxの個数をdp[i][0]、i番目のyがTrueのときのxの個数をdp[i][1]で管理すると、

• i個目のSが"AND"のとき
  yiがTrueなら、x, yi-1もTrue → $dp[i][1] = dp[i-1][1]$
  yiがFalseのとき、yi-1がFalseならxはどちらでも、yi-1がTrueならxはFalse
  → $dp[i][0] = dp[i-1][0]*2+dp[i-1][1] $

• i個目のSが"OR"のとき
  yiがFalseなら、x, yi-1のどちらもFalse → $dp[i][0] = dp[i-1][0]$
  yiがTrueのとき、yi-1がFalseならxはTrue、yi-1がTrueならxはどちらでもOK
  → $dp[i][1] = dp[i-1][0]+dp[i-1][1]*2 $

と状態を定められるので、あとはこれを実装してやれば答えが導ける。

実装

sample.py

N = int(input())

dp = [[0, 0] for _ in range(N+1)]

dp[0][0] = 1
dp[0][1] = 1

for i in range(N):
    s = input()
    if s == "AND":
        dp[i+1][1] = dp[i][1]
        dp[i+1][0] = dp[i][0]*2 + dp[i][1]
    else:
        dp[i+1][1] = dp[i][1]*2 + dp[i][0]
        dp[i+1][0] = dp[i][0]

print(dp[N][1])

最後に

同じ悩みを抱えている方の少しでも力になれれば幸いです。
また、最近、ブログを開設したので、良ければこちらにも遊びに来てください。