ABC267 B - Split? Python解答例

AtCoder Beginner Contest 267 B - Split?をPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

ボウリングのピンは$1$から$10$の番号が付けられており、上から見ると下図のように配置されます。

この図の二つの点線に挟まれた部分を列と呼ぶことにします。
例えば、ピン$1, 5$とピン$3, 9$はそれぞれ同じ列に存在します。

いくつかのピンが倒れた状態のうち、特殊なものはスプリットと呼ばれます。
ピンの配置がスプリットであるとは、以下の条件が全て成り立つことを言います。

• ピン$1$が倒れている。
• ある二つの異なる列であって、次の条件を満たすものが存在する。
  • それぞれの列には、立っているピンが$1$本以上存在する。
  • それらの列の間に、ピンが全て倒れている列が存在する。

具体例は入出力例を参考にしてください。

さて、あるピンの配置が長さ$10$の文字列$S$として与えられます。$i = 1, \dots, 10$について、ピン$i$が倒れているとき$S$の$i$文字目は0であり、ピン$i$が立っているとき$S$の$i$文字目は1です。
$S$で表されるピンの配置がスプリットかどうか判定してください。

制約

• $S$は0と1からなる長さ$10$の文字列

解答例

右端と左端を全探索

まず、7つの列について立っているピンの数を記録します。列の番号を左から$0, 1, \ldots, 6$とします。
長さ7の配列$P$を用意し、$P[i]$をその列に立っているピンの数とします。列$i$に属するピンが立っていたとき、$P[i]$に1を加算していくことでピンの数を記録できます。

次に、右端と左端を全探索します。右端の列の番号を$l$、左端の列の番号を$r$とします。
$l < k < r$を満たす整数$k$について、以下の3つをそれぞれチェックします。

1. 列$l$には立っているピンが1本以上ある
2. 列$r$には立っているピンが1本以上ある
3. ピンが全て倒れている列$k$が存在する

これらの条件を全て満たす$l$, $r$, $k$が存在し、かつピン1が倒れているとき、答えはYesです。

実装例

Pythonによる実装例を以下に示します。Pythonだとitertools.combinationsを用いると組み合わせの全探索が簡単にできます。

b.py

from typing import *
import itertools


def main():
    # 入力受け取り
    S: str = input()

    # ピンの番号から、そのピンが属する列の番号へ変換する辞書
    D: Dict[int, int] = {
        0: 3,
        1: 2,
        2: 4,
        3: 1,
        4: 3,
        5: 5,
        6: 0,
        7: 2,
        8: 4,
        9: 6,
    }

    # 各列に立っているピンの本数を記録する配列
    C: List[int] = [0] * 7
    for i, s in enumerate(S):
        if s == "1":
            C[D[i]] += 1

    # 右端と左端を全探索する
    # 条件を満たすl, r, kが存在するかどうかを表すフラグ
    ok: bool = False
    for l, r in itertools.combinations(range(7), r=2):
        # まず、列Lと列rには立っているピンが1本以上存在しなければならない
        if C[l] >= 1 and C[r] >= 1:
            # lとrの間にある全ての列をチェックする
            for k in range(l + 1, r):
                # ピンが全て倒れている列があれば条件を満たしている
                if C[k] == 0:
                    ok = True

    # 全ての条件を満たし、かつピン1が倒れていれば答えはYes
    # そうでなければ答えはNo
    if ok and S[0] == "0":
        print("Yes")
    else:
        print("No")


if __name__ == "__main__":
    main()

実際の提出はこちら。