ARC 005 C 器物損壊！高橋君 - いろんなことで四苦Hack（仮）

続き。

dijkstraでどうやってやんだよ...って思って、dijkstraの解答を探した。

ferin-tech.hatenablog.com

拡張ダイクストラとは？？？？？？？？？？？？？？？？？？？？？

明日一通り調べよう...

H, W = map(int, input().split())

field = []
start = ()
goal = ()
for i in range(H):
    line = list(input())
    for j, c in enumerate(line):
        if c == 's':
            start = (i, j)
            line[j] = '.'
        elif c == 'g':
            goal = (i, j)
            line[j] = '.'
    field.append(line)

import queue
push = queue.heappush
pop = queue.heappop
INF = 10**9

DYDX = [(0, 1), (1, 0), (0, -1), (-1, 0)]

def dijkstra(s):
    q = [] # リストをぶち込むキュー
    d = [[[INF for k in range(3)] for j in range(W)] for i in range(H)]
    for i in range(H):
        for j in range(W):
            for k in range(3):
                d[i][j][k] = INF
    d[s[0]][s[1]][2] = 0
    q.append((0, s[0], s[1], 2))

    while q:
        tmp = pop(q)
        y = tmp[1]; x = tmp[2]; num = tmp[3]
        # ゴールだった
        if y == goal[0] and x == goal[1]:
            break
        for dy, dx in DYDX:
            ny = y + dy; nx = x + dx
            if 0 <= ny < H and 0 <= nx < W:
                if field[ny][nx] == '.' and d[ny][nx][num] > d[y][x][num] + 1:
                    d[ny][nx][num] = d[y][x][num]+1
                    push(q, (d[ny][nx][num], ny, nx, num))
                if field[ny][nx] == '#' and num >= 1 and d[ny][nx][num-1] > d[y][x][num]+1:
                    d[ny][nx][num-1] = d[y][x][num] + 1
                    push(q, (d[ny][nx][num-1], ny, nx, num-1))
    return d

d = dijkstra(start)

for i in range(3):
    if d[goal[0]][goal[1]][i] != INF:
        print('YES')
        break
else:
    print('NO')

いくつかのテストケースでTLEになった。書き方がちょっと悪いかもな。C++ のvectorの'<'は、lexicographical_compare - cpprefjp C++日本語リファレンスのようだ。cpprefに書いてあった。難しいなぁと思った。もうちょい掘り下げて理解したい。