PythonでJOI難易度5を埋める #22

E - チーズ (Cheese)

考えたこと

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

とあるので，チーズは1からNまでを順番に食べていくことになります．二点の最短距離を求めればよいので，BFSに突っ込みます．

from collections import deque

h, w, n = map(int, input().split())
maps = [input() for _ in range(h)]

def bfs(q, d, maps, h , w):
    s_x = q[0][1]
    s_y = q[0][0]

    while len(q) != 0:
        now = q.popleft()
        for i, j in [[-1, 0], [0, 1], [0, -1], [1 ,0]]:
            go = [now[0] + i, now[1] + j]
            if go[0] < 0 or go[0] >= h or go[1] < 0 or go[1] >= w:
                continue
            if maps[go[0]][go[1]] == 'X':
                continue
            if go[0] == s_y and go[1] == s_x:
                continue
            if d[go[0]][go[1]] == 0:
                d[go[0]][go[1]] = d[now[0]][now[1]] + 1
                q.append([go[0], go[1]])
    return d

#cheese[i][0]はy座標，cheese[i][1]はx座標
cheeses = [None] * (n + 1)
for i in range(h):
    for j in range(w):
        if maps[i][j] != '.' and maps[i][j] != 'X' and maps[i][j] != 'S':
            cheeses[int(maps[i][j])] = [i, j]
        if maps[i][j] == 'S':
            cheeses[0] = [i, j]

ans = 0
for i in range(n):
    q = deque([cheeses[i]])
    d = [[0] * w for _ in range(h)]
    ans += bfs(q, d, maps, h, w)[cheeses[i + 1][0]][cheeses[i + 1][1]]

print(ans)

余談

Pythonは遅い(高速化できてない)，PyPyはメモリを食う(メモリ減らせてない)...