atcoder E - チーズ (Cheese) のpythonでの解答

atcoder.jp

・注意点：TLEする可能性があるのでPypyで提出してください。

・説明：幅優先探索を使ってとく。プログラムの流れはgoalにチーズの位置を入れて、チーズの位置をゴール(gx,gy)、次のループでゴールをスタートにして(sx=gx,sy=gy)幅優先探索をする。また幅優先探索では、for文のはじめで前のループの結果を引き継がないように初期化する必要がある。

・実装：

from collections import deque
H,W,N=map(int,input().split())
field=[]
for _ in range(H):
  field.append(list(input()))

dx=[1,0,-1,0]
dy=[0,1,0,-1]
goal=[0 for _ in range(N)]  #チーズ1，2，3，、、、の座標

#スタートの座標とチーズの座標の特定。
for i in range(H):
  for j in range(W):
    if field[i][j]=='S':
      sx=i
      sy=j
    for k in range(N):
      if field[i][j]==str(k+1):
        goal[k]=(i,j)
#print(goal)

#スタートからの距離がdist
dist=[[-1 for _ in range(W)] for _ in range(H)]
que=deque()

#スタートはdist=0、queにスタートを追加
dist[sx][sy]=0
que.append((sx,sy))

ans=0

#チーズの位置毎に幅優先探索していく
for gx,gy in goal:
  
  #初期化
  dist=[[-1 for _ in range(W)] for _ in range(H)]
  dist[sx][sy]=0
  que.append((sx,sy))
  

  #queがからになるまで幅優先探索をやる
  while len(que)!=0:
    x,y=que.popleft()  #queの先頭のペアを取り出す
  
    #移動
    for i in range(4):
      next_x=x+dx[i]
      next_y=y+dy[i]
    
      #場外はスルー、Xのときも
      if next_x<0 or next_x>=H or next_y<0 or next_y>=W:
        continue
      if field[next_x][next_y]=='X':
        continue
        
      #未探索ならqueにいれる
      if dist[next_x][next_y]==-1:
        que.append((next_x,next_y))
        #距離の更新
        dist[next_x][next_y]=dist[x][y]+1
      
  ans+=dist[gx][gy]  
  #前のゴールからスタートするので
  sx=gx
  sy=gy
  
print(ans)