Сортировка массива по схеме

from random import randint
 
def snake(arr):
    n = len(arr)
    matrix = [[0]*n for _ in range(n)]
    lst = (elem for elem in sorted(sum(arr, [])))
 
    for k in range(2*n):
        tmp = min(k, n-1)
        for j in range(tmp, -1, -1):
            if k%2:
                row = k-j
                col = j 
            else:
                row = j 
                col = k-j
            
            if row < n and col < n:
                matrix[n-1-row][col] = next(lst)
 
    return matrix
 
 
n = 5
a = [[randint(10,99) for c in range(n)] for r in range(n)]
for row in a:
    print(*row)
print()
 
 
for row in snake(a):
    print(*row)

arr=[]
count=0
N=7
for i in range(N):
    arr.append([])
    for j in range(N):
        #count+=1
        arr[i].append(count)
 
print(*arr, sep='\n')
 
# Начало взрыва мозга, обратный отсчет
#
i=6
j=0
k=1
z=1
 
counter=0
print(arr[i][j])
 
 
 
# Новая попытка
for n in range(N*2):
    if i>j:
        if i==6 and j==n:
            arr[i][j]=counter
            counter+=1
            continue
        if n%2:
            j+=z
            arr[i][j]=counter
            counter+=1
            z=-1
            g=j
            for m in range(g):#
                i+=z
                j+=z
                arr[i][j]=counter
                counter+=1
        if not n%2:
            i+=z
            arr[i][j]=counter
            counter+=1
            z=1
            g=i
            for m in range(g,N-1):
                i+=z
                j+=z
                arr[i][j]=counter
                counter+=1
 
    else:
        if n%2:
            i-=z
            arr[i][j]=counter
            counter+=1
            z=-1
            g=i
            for m in range(g):
                i+=z
                j+=z
                arr[i][j]=counter
                counter+=1
 
        if not n%2:
            j-=z
            arr[i][j]=counter
            counter+=1
            z=1
            g=j
            for m in range(g, N-1):
                i+=z
                j+=z
                arr[i][j]=counter
                counter+=1
 
print(*arr, sep='\n')