paiza「十億連勝 (paizaランク S 相当) 」をpython3で解いてみた

初めに

paizaの「十億連勝 (paizaランク S 相当)」を、python3で解いてみた。というよりも、解けなかったので、「解説を見てコードを考えた」というのが正しいかも。

方針

解説を、自分なりに解釈したものです。

• 帰納的に、i - 1 番目のステージまでの組み合わせの数から、i 番目のステージまでの組み合わせの数を求める。
• 各ステージ終了時に出現する状態を表すために、
  
  • ステージ終了時点での連勝数 w と過去に X 連勝したことがあるかの真偽値 b の組 (w, b) を利用する。
    
  • ステージ終了時点での各状態の組み合わせの数 c を dict で保存する。
    
  • 組は、dict のキーにするために、hashable な tuple で表す。
    
• 初期状態は、(0, False) で表せて、組み合わせの数は 1 である。
  
• i - 1 番目のステージ終了時の状態の１つが (w, b) で組み合わせの数が c のとき、i 番目のステージ終了時の状態と組み合わせの数は、次のようになる。
  
  • w + A[i] ≦ X のとき、
  • 全勝する場合は 1 通りで、連勝数は w + A[i] に延び、真偽値 b は変わらないから、
    現在のステージ終了時の状態 (w + A[i], b) の組み合わせの数を c だけ増やす。
  • どこかで負ける場合は A[i] 通りあり、連勝数は 0 になり、、真偽値 b は変わらないから、
    現在のステージ終了時の状態 (0, b) の組み合わせの数を c * A[i] だけ増やす。
    
  • w + A[i] > X のとき、
  • X 連勝して次で負ける場合は 1 通りで、連勝数は 0 になり真偽値は True となるから、
    現在のステージ終了時の状態 (0, True) の組み合わせの数を c だけ増やす。
  • X 連勝未満で負ける場合は X - w 通りあり、連勝数は 0 になり、、真偽値 b は変わらないから、
    現在のステージ終了時の状態 (0, b) の組み合わせの数を c * (X - w) だけ増やす。
  • X + 1 連勝以上する場合は、隠し要素の条件を満たさないので無視する。
    
• dict.get(key, default) で、key が辞書にあれば key に対する値を、そうでなければ default を取得する。
  
• 隠し要素の条件を満たすプレイの仕方の総数は、
  真偽値 b が真か、全ステージが終わった時点での連勝数 w が X になっているものの総和を求める。
  
• 下 9 桁の取り出しには、1000000000 で割った余りを用いる。

コード

解説のC++のコードを参考にしました。

N, X = map(int, input().split())
A = [int(input()) for _ in range(N)]

previous = {(0, False): 1}
for i in range(N):
    current = {}
    for state in previous:
        w = state[0]
        b = state[1]
        c = previous[state]
        if w + A[i] <= X:
            current[(w + A[i], b)] = current.get((w + A[i], b), 0) + c
            current[(0, b)] = current.get((0, b), 0) + c * A[i]
        else:
            current[(0, True)] = current.get((0, True), 0) + c
            current[(0, b)] = current.get((0, b), 0) + c * (X - w)

    previous = current.copy()

answer = 0
for state in current:
    if state[1] or state[0] == X:
        answer += current[state]

print(answer % 1000000000)