【AtCoder版！蟻本】AOJ0558 Cheese【幅優先探索】

AtCoder版！蟻本の幅優先探索で類題としてあげられているAOJ0558 CheeseをPythonで解いていきます。

リンク

AOJ0558

AOJ0558 Cheese

問題文

今年も JOI 町のチーズ工場がチーズの生産を始め，ねずみが巣から顔を出した．JOI 町は東西南北に区画整理されていて，各区画は巣，チーズ工場，障害物，空き地のいずれかである．ねずみは巣から出発して全てのチーズ工場を訪れチーズを 1 個ずつ食べる．

この町には，N 個のチーズ工場があり，どの工場も１種類のチーズだけを生産している．チーズの硬さは工場によって異なっており，硬さ 1 から N までのチーズを生産するチーズ工場がちょうど 1 つずつある．

ねずみの最初の体力は 1 であり，チーズを 1 個食べるごとに体力が 1 増える．ただし，ねずみは自分の体力よりも硬いチーズを食べることはできない．

ねずみは，東西南北に隣り合う区画に 1 分で移動することができるが，障害物の区画には入ることができない．チーズ工場をチーズを食べずに通り過ぎることもできる．すべてのチーズを食べ終えるまでにかかる最短時間を求めるプログラムを書け．ただし，ねずみがチーズを食べるのにかかる時間は無視できる．

入力

入力は H+1 行ある．1 行目には 3 つの整数 H，W，N (1 ≤ H ≤ 1000，1 ≤ W ≤ 1000，1 ≤ N ≤ 9) がこの順に空白で区切られて書かれている．2 行目から H+1 行目までの各行には，'S'，'1', '2', ..., '9'，'X'，'.' からなる W 文字の文字列が書かれており，各々が各区画の状態を表している．北から i 番目，西から j 番目の区画を (i,j) と記述することにすると (1 ≤ i ≤ H, 1 ≤ j ≤ W)，第 i+1 行目の j 番目の文字は，区画 (i,j) が巣である場合は 'S' となり，障害物である場合は 'X' となり，空き地である場合は '.' となり，硬さ 1, 2, ..., 9 のチーズを生産する工場である場合はそれぞれ '1', '2', ..., '9' となる．入力には巣と硬さ 1, 2, ..., N のチーズを生産する工場がそれぞれ 1 つずつある．他のマスは障害物または空き地であることが保証されている．ねずみは全てのチーズを食べられることが保証されている．

解答例

from collections import deque
import copy

# 入力
h, w, n = map(int,input().split())
maze = [] # 迷路の文字配列
for _ in range(h):
    maze.append(input())

# 各座標への最短距離の配列（INFで初期化）
INF = int(1e9)
D = [[INF] * w for _ in range(h)]

# 4方向の移動ベクトル
dx = [1,0,-1,0]
dy = [0,1,0,-1]

# スタート位置の探索
for i in range(h):
    for j in range(w):
        if maze[i][j] == "S":
            sy = i
            sx = j


# 硬さxのチーズ位置の探索
def search_cheese(x):
    for i in range(h):
        for j in range(w):
            if maze[i][j] == str(x):
                gy = i
                gx = j
    return gy, gx # 硬さxのチーズ座標を返す


# スタートからゴールへの最短距離を求めるbfs
def bfs():
    que = deque()
    que.append([sy,sx]) # スタート地点をキューに入れる
    D_copy[sy][sx] = 0  # スタート地点の距離を0とする

    # キューが空になるまで探索
    while len(que) != 0:
        p = que.popleft()   # 先入れのキューの要素を取り出す
        if p[0] == gy and p[1] == gx:   # ゴールなら探索終了
            break
        for i in range(4):  # pから4方向に移動 移動後の点は(nx,ny)
            ny = p[0] + dy[i]
            nx = p[1] + dx[i]
            # 移動可能か、この探索で訪れたことがあるかの判定
            if 0 <= nx < w and 0 <= ny < h and maze[ny][nx] != "X" and D_copy[ny][nx] == INF:
                # 移動可能ならキューに入れて、その点の距離をpからの距離+1で確定する
                que.append([ny,nx])
                D_copy[ny][nx] = D_copy[p[0]][p[1]] + 1
    # ゴールまでの最短経路を返す
    return D_copy[gy][gx] 
                


# 硬さNまでの最短経路探索
res = 0 # 最短経路
for k in range(n):
    D_copy = copy.deepcopy(D)
    tmp = search_cheese(k+1)
    gy = tmp[0]
    gx = tmp[1]
    res += bfs()
    sy = gy
    sx = gx

print(res)

解説

N個のチーズ工場があり、それぞれ1からNまでの硬さのチーズを生産しています。
スタート地点から硬さNのチーズを食べるまでの最短経路を答える問題です。

この問題も幅優先探索（BFS)で最短経路を求めていきます。

入力を受け取ったあと、BFSの準備として迷路と同じ大きさの配列D（INFで初期化）を用意します。
迷路のスタート地点はSで表されているので、全探索してxy座標をsx,syに代入しました。

硬さiのチーズの位置はプログラムで何度も探す必要があるのでsearch_cheese関数を作成しました。
iを引数として渡すことで硬さiのチーズの位置をgx,gyとして返します。

BFSはスタート地点を渡すことで、gx,gy地点までの最短経路を返す関数です。
最短経路のリストDは以前の探索のデータを引き継いだら間違った答えになりますので、deepcopyしたD_copyに最短経路を記録していきました。

最後に1からNまでのチーズの最短経路をfor文で求めます。
次の目的地であるk+1の硬さのチーズの位置をsearch_sheese関数で求め、BFSで現在地から求めた地点までの最短経路をresに足していきました。
BFSが終わるとゴール地点を次のスタート地点として、次回のループに備えます。

ループを抜ければ求める最短経路がresとなりますので出力すればOKでした。

AtCoder版！蟻本（初級編）の記事リンクページはこちら
ebisuke33.hatenablog.com

リンク