Problem statement

https://leetcode.com/problems/reshape-the-matrix/

Solution

First approach is to reshape to line and then reshape it to new shape. However we can do smarter: just iterate over all elements line by line and use res[count//c][count\%c] = nums[i][j] to fill element by element.

Complexity

Time complexity is O(mn), space complexity is O(mn), but in fact it is O(1) if we do not count output array.

Code

class Solution:
    def matrixReshape(self, nums, r, c):
        m, n, count = len(nums), len(nums[0]), 0
        if m*n != r*c: return nums
        res = [[0] * c for _ in range(r)]
        for i, j in product(range(m), range(n)):
            res[count//c][count%c] = nums[i][j]
            count += 1      
        return res