Заполнение матрицы по диагоналям

Заполнить матрицу числами по порядку по диагоналям, параллельным побочной. Начало — верхний левый угол. Размер матрицы вводится пользователем.

Образец:

n = int(input())
a = [[0] * n for i in range(n)]
k = 1

for i in range(n):  # заполнение выше побочной диагонали
    for j in range(i + 1):
        a[i - j][j] = str(k).ljust(3)
        k += 1
for i in range(n - 1, 0, -1):  # нижняя часть
    for j in range(i):
        a[n - j - 1][n - i + j] = str(k).ljust(3)
        k += 1
        
for row in a:
    print(*row)