蛇形矩阵

def solve(n):
    matrix = [[0] * n for _ in range(n)]
    num = 1
    for diag in range(2 * n - 1):
        if diag < n:
            row, col = 0, diag
            while col >= 0:
                matrix[row][col] = num
                num += 1
                row += 1
                col -= 1
        else:
            row, col = diag - n + 1, n - 1
            while row < n:
                matrix[row][col] = num
                num += 1
                row += 1
                col -= 1
    for row in matrix:
        print(' '.join(map(str, row)))