ABC201のC問題 Secret Numberを解く

ABC201のC問題を解説します。

atcoder.jp

問題

高橋くんは、暗証番号を忘れてしまいました。暗証番号は 0 から 9 までの数字のみからなる 4 桁の文字列で、0 から始まる場合もあります。
0 から 9 までの各数字について、高橋くんは以下のように記憶しています。彼の記憶は長さ 10 の文字列 によって表されます。

が o のとき : 数字 i は暗証番号に確実に含まれていた。
が x のとき : 数字 i は暗証番号に確実に含まれていなかった。
が ? のとき : 数字 i が暗証番号に含まれているか分からない。
高橋くんが忘れてしまった暗証番号としてあり得るものは何通りありますか？

難しいポイント

3桁以下の数字は頭に0がつくのは、どう処理しようか悩みました。
文字列で数字が含まれているか判定しようか、各桁の数字を抜き出して数字として判定しようか、地味に迷いました。

解き方の方針

4桁なので、0から9999まで全量を探索して、条件に合ってるものをカウントします。

実装

数字は文字列に変換して、oがついた数字が全て含まれているか、xがついた数字が含まれていないことを確認します。
3桁以下の数字については頭に0がつくことが確定なので、0がxの場合、暗証番号の可能性は無くなります。
0がoの時はスルーして次の文字を検索にまわします。

判定する関数は以下のようになりました。

def is_num_ok(num: str) -> bool:
    for o in ok_num:
        if o == '0' and len(num) <= 3:
            continue
        if not(o in num):
            return False
    for n in ng_num:
        if n == '0' and len(num) <= 3:
            return False
        if n in num:
            return False
    return True

最終的なコードは以下のようになりました。

S = input()
ok_num = []
ng_num = []

count = 0

for i, s in enumerate(S):
    if s == 'o':
        ok_num.append(str(i))
    if s == 'x':
        ng_num.append(str(i))


def is_num_ok(num: str) -> bool:
    for o in ok_num:
        if o == '0' and len(num) <= 3:
            continue
        if not(o in num):
            return False
    for n in ng_num:
        if n == '0' and len(num) <= 3:
            return False
        if n in num:
            return False
    return True


for j in range(10000):
    str_j = str(j)
    if is_num_ok(str_j):
        count += 1

print(count)

意外と方針はすぐに思いつきましたが、もっと早く解けるようになりたいと思いました。