Problem p02670 (python)

Fixed submission:
import os

import itertools
import sys

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7


# MOD = 998244353


def test():
    ans_hist = [INF] * (N ** 2)

    ans = 0
    dist = [INF] * (N ** 2)
    for i, p in enumerate(P):
        p -= 1
        ph, pw = p // N, p % N
        d = min(ph, pw, N - ph - 1, N - pw - 1)

        j = i - 1
        while j >= 0:
            q = P[j] - 1
            qh, qw = q // N, q % N
            d = min(d, dist[q] + abs(ph - qh) + abs(pw - qw) - 1)
            j -= 1

        dist[p] = d
        ans_hist[p] = d

        j = i - 1
        while j >= 0:
            q = P[j] - 1
            qh, qw = q // N, q % N
            dist[q] = min(dist[q], dist[p] + abs(ph - qh) + abs(pw - qw) - 1)
            j -= 1

        ans += d
        # print(p, dist)
    print(ans)
    # print(dist)
    # print(np.array(ans_hist, dtype=int).reshape(N, N))
    return ans


N = int(sys.stdin.buffer.readline())
P = list(map(int, sys.stdin.buffer.readline().split()))

# 解説
dist = [INF] * len(P)
for h, w in itertools.product(range(N), range(N)):
    dist[h * N + w] = min(h, w, N - h - 1, N - w - 1)
present = [1] * len(P)

ans = 0
for p in P:
    p -= 1
    h, w = p // N, p % N
    present[p] = 0

    ans += dist[p]
    stack = [(p, dist[p])]
    while stack:
        p, d = stack.pop()
        h, w = p // N, p % N
        d += present[p]
        for dh, dw in ((0, 1), (1, 0), (0, -1), (-1, 0)):
            if 0 <= h + dh < N and 0 <= w + dw < N:
                q = (h + dh) * N + w + dw
                if d < dist[q]:
                    dist[q] = d
                    stack.append((q, d))
print(ans)
# test()

Buggy submission (verdict: Runtime Error):
import os

import itertools
import sys

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7


# MOD = 998244353


def test():
    ans_hist = [INF] * (N ** 2)

    ans = 0
    dist = [INF] * (N ** 2)
    for i, p in enumerate(P):
        p -= 1
        ph, pw = p // N, p % N
        d = min(ph, pw, N - ph - 1, N - pw - 1)

        j = i - 1
        while j >= 0:
            q = P[j] - 1
            qh, qw = q // N, q % N
            d = min(d, dist[q] + abs(ph - qh) + abs(pw - qw) - 1)
            j -= 1

        dist[p] = d
        ans_hist[p] = d

        j = i - 1
        while j >= 0:
            q = P[j] - 1
            qh, qw = q // N, q % N
            dist[q] = min(dist[q], dist[p] + abs(ph - qh) + abs(pw - qw) - 1)
            j -= 1

        ans += d
        # print(p, dist)
    print(ans)
    # print(dist)
    # print(np.array(ans_hist, dtype=int).reshape(N, N))
    return ans


N = int(sys.stdin.buffer.readline())
P = list(map(int, sys.stdin.buffer.readline().split()))

# 解説
dist = [[INF] * N for _ in range(N)]
for h, w in itertools.product(range(N), range(N)):
    dist[h][w] = min(h, w, N - h - 1, N - w - 1)
present = [True] * (N ** N)

ans = 0
for p in P:
    p -= 1
    h, w = p // N, p % N
    present[p] = False

    ans += dist[h][w]
    stack = [(h, w, dist[h][w])]
    while stack:
        h, w, d = stack.pop()
        d += present[h * N + w]
        for dh, dw in ((0, 1), (1, 0), (0, -1), (-1, 0)):
            if 0 <= h + dh < N and 0 <= w + dw < N:
                if d < dist[h + dh][w + dw]:
                    dist[h + dh][w + dw] = d
                    stack.append((h + dh, w + dw, d))
print(ans)
# test()