566. Reshape the Matrix

class Solution {
public:
    //return the n-th element of m using row traversing
    int rowTraverse(vector<vector<int>>& m, int n) {
        for (vector<int>& row : m) {
            if (n < row.size()) {
                return row[n];
            }
            n -= row.size();
            continue;
        }
    }
    
    vector<vector<int>> matrixReshape(vector<vector<int>>& nums, int r, int c) {
        //check the invalid cases
        //since its the matrix, we can assume that every row has the same amount of elements.
        int rows = nums.size();
        int cols = nums[0].size();
        if (r*c != rows*cols) {
            //is that the only case?
            //invalid reshape, return original matrix
            return nums;
        }
        
        vector<vector<int>> m;
        for (int i = 0; i < r; i++) {
            //fill the row
            vector<int> _r;
            for (int j = 0; j < c; j++) {
                _r.emplace_back(rowTraverse(nums, i*c+j));
            }
            m.emplace_back(_r);
        }
        return m;
    }
};