Заполнить двумерный массив

def gen_array(n, m):
    matrix = [[0]*m for _ in range(n)]
    num = 1
    for k in range(n + m):
        for j in range(min(k, m - 1), -1, -1):
            i = k - j
            if i < n:
                matrix[i][j] = num 
                num += 1
    return matrix
 
 
n = 5
arr = gen_array(n, n)
k = len(str(arr[-1][-1]))
for row in arr:
    print(' '.join([f'{elem:{k}d}' for elem in row]))

def gen_array(n, m):
    matrix = [[0]*m for _ in range(n)]
    num = 1
 
    for k in range(n + m):
        tmp = min(k, max(n-1, m-1))
        for j in range(tmp, -1, -1):
            if k%2:
                row = j
                col = k - j 
            else:
                row = k - j 
                col = j
            
            if row < n and col < m:
                matrix[row][col] = num
                num += 1
 
    return matrix
 
n = int(input())
arr = gen_array(n, n) 
for row in arr:
    print(' '.join([f'{elem:{k}d}' for elem in row]))