Reshape the Matrix

LeetCode題目: 566. Reshape the Matrix

My solution:

/**
 * @param {number[][]} mat
 * @param {number} r
 * @param {number} c
 * @return {number[][]}
 */
let matrixReshape = (mat, r, c) => {
    let result = [], tmpArr = [];
    let flatMat = mat.flat();
    if(flatMat.length === r * c) {
        for(let i=0; i<r; i++) {
            for(let j=0; j<c; j++){
                tmpArr.push(flatMat.shift());
            }
            result.push(tmpArr);
            tmpArr = [];
        }
        return result
    } else return mat;
};