ebisukeプログラミング初心者脱出黙示録

30歳を過ぎてから始めたプログラミングと競プロの記録。Pythonで取り組んでいます。Arduinoで電子工作も

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

f:id:ebisuke33:20210317144704p:plain
AtCoder版!蟻本幅優先探索で類題としてあげられているAOJ0558 CheeseをPythonで解いていきます。

AOJ0558

AOJ0558 Cheese

問題文

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

この町には,N 個のチーズ工場があり,どの工場も1種類のチーズだけを生産している.チーズの硬さは工場によって異なっており,硬さ 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