Problem p03426 (python)

Fixed submission:
import sys

import numba as nb
import numpy as np

input = sys.stdin.readline


@nb.njit("void(i8,i8,i8,i8[:,:],i8,i8[:,:],i8[:])", cache=True)
def dfs(H, W, D, A, Q, LR, ans):
    position = [(-1, -1)] * (H * W + 1)
    for h in range(H):
        for w in range(W):
            position[A[h][w]] = (h, w)

    mp = np.zeros(shape=(H * W + 1), dtype=np.int64)

    for i in range(1, H * W + 1 - D):
        h, w = position[i]
        y, x = position[i + D]
        mp[i + D] = abs(x - w) + abs(y - h)

    for d in range(1, D + 1):
        mp[d::D] = np.cumsum(mp[d::D])

    for i in range(Q):
        l, r = LR[i]
        ans[i] = mp[r] - mp[l]


def main():
    H, W, D = map(int, input().split())
    A = np.zeros(shape=(H, W), dtype=np.int64)
    for i in range(H):
        A[i] = tuple(map(int, input().split()))
    Q = int(input())
    LR = np.zeros(shape=(Q, 2), dtype=np.int64)
    for i in range(Q):
        LR[i] = tuple(map(int, input().split()))

    ans = np.zeros(shape=Q, dtype=np.int64)
    dfs(H, W, D, A, Q, LR, ans)

    print("\n".join(map(str, ans)))


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import sys

import numba as nb
import numpy as np

input = sys.stdin.readline


@nb.njit("void(i8,i8,i8,i8[:,:],i8[:,:])", cache=True)
def dfs(H, W, D, A, mp):
    position = [(-1, -1)] * (H * W + 1)
    for h in range(H):
        for w in range(W):
            position[A[h][w]] = (h, w)

    for i in range(1, H * W + 1 - D):
        h, w = position[i]
        y, x = position[i + D]
        mp[i][i + D] = abs(x - w) + abs(y - h)

    for i in range(1, H * W + 1 - 2 * D):
        for j in range(i + 2 * D, H * W + 1, D):
            mp[i][j] = mp[i][j - D] + mp[j - D][j]


def main():
    H, W, D = map(int, input().split())
    A = np.zeros(shape=(H, W), dtype=np.int64)
    for i in range(H):
        A[i] = list(map(int, input().split()))
    Q = int(input())
    L = [0] * Q
    R = [0] * Q
    for i in range(Q):
        L[i], R[i] = map(int, input().split())

    mp = np.zeros(shape=(H * W + 1, H * W + 1), dtype=np.int64)
    dfs(H, W, D, A, mp)

    ans = [mp[l, r] for l, r in zip(L, R)]
    print("\n".join(map(str, ans)))


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import sys

import numba as nb
import numpy as np

input = sys.stdin.readline


@nb.njit("void(i8,i8,i8,i8[:,:],i8[:,:])", cache=True)
def dfs(H, W, D, A, mp):
    position = [(-1, -1)] * (H * W + 1)
    for h in range(H):
        for w in range(W):
            position[A[h][w]] = (h, w)

    for i in range(1, H * W + 1 - D):
        h, w = position[i]
        y, x = position[i + D]
        mp[i][i + D] = abs(x - w) + abs(y - h)

    for i in range(1, H * W + 1 - 2 * D):
        for j in range(i + 2 * D, H * W + 1):
            mp[i][j] = mp[i][j - D] + mp[j - D][j]


def main():
    H, W, D = map(int, input().split())
    A = np.zeros(shape=(H, W), dtype=np.int64)
    for i in range(H):
        A[i] = list(map(int, input().split()))
    Q = int(input())
    L = [0] * Q
    R = [0] * Q
    for i in range(Q):
        L[i], R[i] = map(int, input().split())

    mp = np.zeros(shape=(H * W + 1, H * W + 1), dtype=np.int64)
    dfs(H, W, D, A, mp)

    ans = [mp[l, r] for l, r in zip(L, R)]
    print("\n".join(map(str, ans)))


if __name__ == "__main__":
    main()