Pythonで毎日AtCoder #31

はじめに

前回
#24の続きです。今日はARC031-Bを解きます。

#31

ARC031-B

考えたこと
最初はDFSの書きかたが分からなかったので、記事を読みました。しかし、検索しても再帰関数を使ったものしか出てこなかったので、今回はstackを使って解きます(それしか知らない)。
この問題の一つ前の問題との違いはゴールが設定されていないことです。最初はこれで躓きました。ゴールが無いので探索の終了条件をうまく設定できなかったからです。ですので、ゴールを考えます。この問題は'o'が全て繋がっているかどうかを聞いている問題なので、探索した'o'数が入力の'o'の合計になることがゴールです。また、埋め立ては二重forで計算しています。変数名が汚ないのは許してください。

from collections import deque
field = [list(input()) for _ in range(10)] #strで受け取ると埋め立て時にエラー吐くのでlistにしてます

def dfs(field,visited_list,stack):
    step = 0
    while len(stack) > 0:
        h, w = stack.pop()
        for j, k in ([1,0],[-1,0],[0,1],[0,-1]):
            new_h, new_w = h+j, w+k
            if new_h < 0 or new_h >= 10 or new_w < 0 or new_w >= 10:
                continue
            if field[new_h][new_w] != 'x' and visited_list[new_h][new_w] == 0:
                visited_list[new_h][new_w] = 1
                stack.append([new_h,new_w])
                step += 1 #ステップを計算
    return step

count_step = 0
for i in range(10):
    for j in range(10):
        if field[i][j] == 'o':
            count_step += 1
ume = 0
for i in range(10):
    for j in range(10):
        if field[i][j] != 'o':
            field[i][j] = 'o'
            count_step += 1 #埋め立てた分'o'が増える
            ume = True
        visited_list = [[0]*10 for _ in range(10)]
        visited_list[i][j] = 1
        stack = deque([[i,j]])
        step = dfs(field,visited_list,stack)
        step += 1
        if step == count_step:
            print('YES')
            quit()
        if ume:
            field[i][j] = 'x' #埋め立て後は戻す
            ume = False
            count_step -= 1
print('NO')

まとめ

少しずつ慣れていって使いこなせるようになりたい。おもしろい問題があったらリプください。ではまた、おやすみなさい。