0
votes

I need create a Matlab mex function that will take an input matrix and return the matrix diagonal for ex.
input = ([1,2,3;4,5,6])
output =
1 2 3 0 0 0
0 0 0 4 5 6

My problem is since Matlab reads matrixes columns wise not row wise my mex function for this input will output
1 4 0 0 0 0
0 0 2 5 0 0
0 0 0 0 3 6

My code is as follows:

#include <matrix.h>
#include <mex.h>

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
        mxArray *a_in, *b_out;
        const mwSize *dims;
        double *a, *b; 
        int rows, cols;

        // Input array:
        a_in = mxDuplicateArray(prhs[0]);

        // Get dimensions of input array:
        dims = mxGetDimensions(prhs[0]);
        rows = (int) dims[0];
        cols = (int) dims[1];

        // Output array:
        if(rows == cols){
            b_out = plhs[0] = mxCreateDoubleMatrix(rows, rows*cols, mxREAL);
        }
        else{
            b_out = plhs[0] = mxCreateDoubleMatrix(cols, rows*cols, mxREAL);
        }
        // Access the contents of the input and output arrays:
        a = mxGetPr(a_in);
        b = mxGetPr(b_out);
        

        // Compute exdiag function of the input array
        int count = 0;
        for (int i = 0; i < rows; i++) {
            for(int j = 0; j<cols;j++){
                if(rows == cols){
                    b[rows*count+count/rows] = a[j + rows * i];
                    count++;
                }
                else if(rows < cols){
                    b[cols*count+count/rows] = a[j + cols * i];
                    count++;
                }
                else if(rows>cols){
                    b[cols*count+count/rows] = a[j + cols * i];
                    count++;
                }
            }
    }
}

How would you go about changing my code to read the input matrix row wise so I can have the proper output?