PythonでJOI難易度5を埋める #10

A - 碁石ならべ

考えたこと
石の状態をlistで管理します。配列に(石の色,連続する個数)のtupleを保存して操作します。問題文と同じく白は0、黒は1で保存します。
例:([(0, 2), (1, 2), (0, 3)])これは、0011000を表します

まずは、iが偶数のときを考えます。iが偶数のときは、後ろに石を追加するだけで以前の配置を変更することはありません。

  • 追加する石が最後の石と同じ色の場合は、\cdots ,(色,個数) → \cdots ,(色,個数+1)
  • 追加する石が最後の石と異なる色の場合は、\cdots ,(色,個数) → \cdots ,(色,個数) ,(追加した石の色,1)

それぞれ、popとappendで実装できます。

次にiが奇数のときを考えます。iが奇数のときは、後ろに石を追加する操作とオセロのように挟まれた石の色を変更する操作が必要です。

  • 追加する石が最後の石と同じ色の場合は、\cdots ,(色,個数) → \cdots ,(色,個数+1) ←iが偶数で追加する石が同じ色の場合と一緒
  • 追加する石が最後の石と異なる色の場合は、後ろの二つを見る必要があります。\cdots ,(色1,個数1), (色2,個数2) → \cdots ,(色1,個数1 + 個数2 + 1)

になります。この操作もpopとappendで実装できます。 あとは、0の数を数えるだけです。

n = int(input())
stones = [int(input()) for _ in range(n)]

state = []
if stones[0] == 0:
    state.append((0, 1))
else:
    state.append((1, 1))

for i in range(1,n):
    m1 = state.pop(-1)
    if (i + 1) % 2 == 0:
        if m1[0] == stones[i]:
            state.append((stones[i], m1[1] + 1))
        else:
            if len(state) != 0:
                m2 = state.pop(-1)
                state.append((stones[i], m2[1] + m1[1] + 1))
            else:
                state.append((stones[i], m1[1] + 1))
    else:
        if m1[0] == stones[i]:
            state.append((stones[i], m1[1] + 1))
        else:
            state.append(m1)
            state.append((stones[i], 1))

ans = 0
for i in range(len(state)):
    if state[i][0] == 0:
        ans += state[i][1]
print(ans)

PyPyだとメモリが厳しいのでPythonで提出しましょう。