566. Reshape the Matrix

Python

• Time: O(m*n)
• Space: O(1)

class Solution:
    def matrixReshape(self, mat: List[List[int]], r: int, c: int) -> List[List[int]]:
        new_mat = []
        new_row = []

        for row in mat:
            for val in row:
                new_row.append(val)
                if len(new_row) >= c:
                    new_mat.append(new_row)
                    new_row = []

        return new_mat if len(new_mat) == r and len(new_mat[-1]) == c else mat