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):
from itertools import product


def solve(string):
    h, w, d, *a = map(int, string.split())
    a, q, lr = a[:h * w], a[h * w], a[h * w + 1:]
    n = h * w
    index = [a.index(i + 1) for i in range(h * w)]
    costs = [
        [0 if j <= i or (j - i) % d != 0 else n * (j - i) // d for j in range(n)] for i in range(n)
    ]
    for s, g in product(range(n), repeat=2):
        if costs[s][g] == 0:
            continue
        if s - d >= 0:
            costs[s][g] = costs[s - d][g] - costs[s - d][s]
            continue
        before = 0 if g - d < 0 else costs[s][g - d]
        costs[s][g] = abs(index[g - d] // w - index[g] // w) + abs(index[g - d] % w -
                                                                   index[g] % w) + before
    print(costs)
    return "\n".join([str(costs[l - 1][r - 1]) for l, r in zip(lr[::2], lr[1::2])])


if __name__ == '__main__':
    n, m, d = map(int, input().split())
    ins = '{} {} {}\n'.format(n, m, d) + '\n'.join([input() for _ in range(n)])
    n = int(input())
    print(solve(ins + '\n{}\n'.format(n) + '\n'.join([input() for _ in range(n)])))


from functools import partial
from itertools import islice


def take(n, iterable):
    return list(islice(iterable, n))


def chunked(iterable, n):
    return iter(partial(take, n, iter(iterable)), [])


def distance(x, y):
    return abs(x[0] - y[0]) + abs(x[1] - y[1])


def solve(s):
    h, w, d, *aqlr = map(int, s.split())
    n = h * w
    pos = [(i // w, i % w) for i, _a in enumerate(aqlr[:n])]
    cost = [0] * (n + 1)
    for i in range(d + 1, n + 1):
        cost[i] = cost[i - d] + distance(pos[i - d], pos[i])
    return "\n".join(["%d" % (cost[r] - cost[l]) for l, r in chunked(aqlr[n + 1:], 2)])


n, m, d = map(int, input().split())
l = '{} {} {}\n'.format(n, m, d) + '\n'.join([input() for _ in range(n)])
q, = map(int, input().split())
l += '\n{}\n'.format(q) + '\n'.join([input() for _ in range(q)])
print(solve(l))


def solve(string):
    h, w, d, *a = map(int, string.split())
    a, q, lr = a[:h * w], a[h * w], a[h * w + 1:]
    n = h * w
    index = [a.index(i + 1) for i in range(h * w)]
    # costs = [
    # [0 if j <= i or (j - i) % d != 0 else n * (j - i) // d for j in range(n)] for i in range(n)
    # ]
    diff = [abs(i0 // w - i1 // w) + abs(i0 % w - i1 % w) for i0, i1 in zip(index, index[d:])]
    """
    for s in range(n):
        if sum(costs[s]) == 0:
            continue
        for g in range(n):
            if costs[s][g] == 0:
                continue
            if s - d >= 0:
                costs[s][g] = costs[s - d][g] - costs[s - d][s]
                continue
            before = 0 if g - d < 0 else costs[s][g - d]
            costs[s][g] = +before
    return "\n".join([str(costs[l - 1][r - 1]) for l, r in zip(lr[::2], lr[1::2])])
    """
    print(diff)
    return "\n".join([str(sum(diff[l - 1:r - 1:d])) for l, r in zip(lr[::2], lr[1::2])])


if __name__ == '__main__':
    n, m, d = map(int, input().split())
    ins = '{} {} {}\n'.format(n, m, d) + '\n'.join([input() for _ in range(n)])
    n = int(input())
    print(solve(ins + '\n{}\n'.format(n) + '\n'.join([input() for _ in range(n)])))


from functools import partial
from itertools import islice


def take(n, iterable):
    return list(islice(iterable, n))


def chunked(iterable, n):
    return iter(partial(take, n, iter(iterable)), [])


def distance(x, y):
    return abs(x[0] - y[0]) + abs(x[1] - y[1])


def solve(s):
    h, w, d, *aqlr = map(int, s.split())
    a, q, lr = aqlr[:h * w], aqlr[h * w], aqlr[h * w + 1:]
    n = h * w
    pos = [0] * (n + 1)
    for i, _a in enumerate(a):
        pos[_a] = (i // w, i % w)
    cost = [0] * (n + 1)
    for i in range(d + 1, n + 1):
        cost[i] = cost[i - d] + distance(pos[i - d], pos[i])
    return "\n".join(["%d" % (cost[r] - cost[l]) for l, r in chunked(lr, 2)])


n, m, d = map(int, input().split())
l = '{} {} {}\n'.format(n, m, d) + '\n'.join([input() for _ in range(n)])
q, = map(int, input().split())
l += '\n{}\n'.format(q) + '\n'.join([input() for _ in range(q)])
print(solve(l))


Verdicts: Wrong Answer

Runtime Error

Wrong Answer

Accepted