Problem:
Tonight, in your favourite cinema they are giving the movie Joker and all seats are occupied. In the cinema there are N rows with N seats each, forming an N\times N square. We denote with 1, 2,\dots, N the viewers in the first row (from left to right); with N+1, \dots, 2N the viewers in the second row (from left to right); and so on until the last row, whose viewers are denoted by N^2-N+1,\dots, N^2.  At the end of the movie, the viewers go out of the cinema in a certain order: the i-th viewer leaving her seat is the one denoted by the number P_i. The viewer P_{i+1} waits until viewer P_i has left the cinema before leaving her seat. To exit from the cinema, a viewer must move from seat to seat until she exits the square of seats (any side of the square is a valid exit). A viewer can move from a seat to one of its 4 adjacent seats (same row or same column). While leaving the cinema, it might be that a certain viewer x goes through a seat **currently** occupied by viewer y; in that case viewer y will hate viewer x forever. Each viewer chooses the way that minimizes the number of viewers that will hate her forever.  Compute the number of pairs of viewers (x, y) such that y will hate x forever.

Incorrect submission (verdict: Time Limit Exceeded):
N=int(input())
S=list(map(int,input().split()))
count=0
P=[0]*(N*N+1)
for j in range(N*N):
  P[j+1]=S.index(j+1)
for i in range(N*N):
  player=S[i]
  hayasa=P[player]
  
  que=[[] for _ in range(N)]
  d=[-1]*(N*N+1)
  que[0].append(player)
  d[player]=0
  
  j=0
  finish=0
  while True:
    for kouho in que[j]:
      if kouho%N==0:
        finish=1
        break
      if kouho%N==1:
        finish=1
        break
      if kouho+N>N*N:
        finish=1
        break
      if kouho<N:
        finish=1
        break
        
      if d[kouho+N]==-1:
        if P[kouho+N]>hayasa:
          que[j+1].append(kouho+N)
          d[kouho+N]=0
        else:
          que[j].append(kouho+N)
          d[kouho+N]=0
          
      if d[kouho-N]==-1:
        if P[kouho-N]>hayasa:
          que[j+1].append(kouho-N)
          d[kouho-N]=0
        else:
          que[j].append(kouho-N)
          d[kouho-N]=0
          
      if d[kouho+1]==-1:
        if P[kouho+1]>hayasa:
          que[j+1].append(kouho+1)
          d[kouho+1]=0
        else:
          que[j].append(kouho+1)
          d[kouho+1]=0
          
      if d[kouho-1]==-1:
        if P[kouho-1]>hayasa:
          que[j+1].append(kouho-1)
          d[kouho-1]=0
        else:
          que[j].append(kouho-1)
          d[kouho-1]=0
    if finish==1:
      count+=j

Fixed, accepted submission:
from numba import njit
import numpy as np

@njit
def main(N, P):
    lista = [1, -1, N, -N]
    d = np.zeros((N, N), dtype=np.int32)
    d = d.ravel()
    sit = np.ones_like(d)
    stack = np.empty_like(d)
    count = 0
    for i in range(N*N):
        d[i] = min(i % N, (N-1-i) % N, i//N, N-1-i//N)
    for player in P:
        sit[player-1] = 0
        count += d[player-1]
        stack[0] = player
        num = 0
        while num >= 0:
            p = stack[num]
            num -= 1
            a = d[p-1]+sit[p-1]
            for i in lista:
                if p+i < 1 or p + i > N*N:
                    continue
                if d[p+i-1] > a:
                    d[p+i-1] = a
                    num += 1
                    stack[num] = p+i
    print(count)

N = int(input())
P = np.array(input().split(), np.int32)

main(N, P)