Problem: We have a grid with H rows and W columns. The square at the i-th row and the j-th column will be called Square (i,j).  The integers from 1 through H×W are written throughout the grid, and the integer written in Square (i,j) is A_{i,j}.  You, a magical girl, can teleport a piece placed on Square (i,j) to Square (x,y) by consuming |x-i|+|y-j| magic points.  You now have to take Q practical tests of your ability as a magical girl.  The i-th test will be conducted as follows:    * Initially, a piece is placed on the square where the integer L_i is written.    * Let x be the integer written in the square occupied by the piece. Repeatedly move the piece to the square where the integer x+D is written, as long as x is not R_i. The test ends when x=R_i.    * Here, it is guaranteed that R_i-L_i is a multiple of D.  For each test, find the sum of magic points consumed during that test.

Submission attempts (in order):
import numpy as np
from numba import njit, i8

h, w, d = map(int, input().split())
arr = np.zeros([h, w], dtype=np.int64)

for i in range(h):
    arr[i] = list(map(int, input().split()))

q = int(input())
lr = np.zeros([q, 2], dtype=np.int64)

for i in range(q):
    lr[i] = list(map(int, input().split()))


@njit
def solve(h, w, d, arr):
    result = [[0] for _ in range(d)]
    for i in range(d):
        l = i
        score = 0
        while l + d <= h * w:
            if l == 0:
                result[i].append(score)
                l += d
                continue
            u = np.where(arr == l)
            v = np.where(arr == l + d)
            score += abs(v[0][0] - u[0][0]) + abs(v[1][0] - u[1][0])
            result[i].append(score)
            l += d
    return result


result = solve(h, w, d, arr)


for l, r in lr:
    print(result[l % d][r // d] - result[l % d][l // d])

import numpy as np

h, w, d = map(int, input().split())
arr = np.zeros([h, w], dtype=np.int64)

for i in range(h):
    arr[i] = list(map(int, input().split()))

q = int(input())
lr = np.zeros([q, 2], dtype=np.int64)

for i in range(q):
    lr[i] = list(map(int, input().split()))


cnt = h * w
result = [[0] for _ in range(d)]
for i in range(d):
    l = i
    score = 0
    while l + d <= cnt:
        if l == 0:
            result[i].append(score)
            l += d
            continue
        u = np.where(arr == l)
        v = np.where(arr == l + d)
        score += abs(v[0][0] - u[0][0]) + abs(v[1][0] - u[1][0])
        result[i].append(score)
        l += d


result = solve(h, w, d, arr)


for l, r in lr:
    ans = result[l % d]
    print(ans[r // d] - ans[l // d])

h, w, d = map(int, input().split())
arr = [0] * (h * w + 1)

a = [list(map(int, input().split())) for _ in range(h)]

for i in range(h):
    for j in range(w):
        s = a[i][j]
        arr[s] = (i, j)

q = int(input())
lr = [list(map(int, input().split())) for _ in range(q)]

cnt = h * w
result = [[0] for _ in range(d)]
for i in range(d):
    l = i
    score = 0
    while l + d <= cnt:
        if l == 0:
            result[i].append(score)
            l += d
            continue
        (ux, uy) = arr[l]
        (vx, vy) = arr[l + d]
        score += (abs(vx - ux) + abs(vy - uy))
        result[i].append(score)
        l += d


for l, r in lr:
    ans = result[l % d]
    print(ans[r // d] - ans[l // d])

Verdicts: Time Limit Exceeded

Runtime Error

Accepted