class Solution:
    def matrixReshape(self, nums: List[List[int]], r: int, c: int) -> List[List[int]]:
        """
        tc O(R*C) sp O(C)
        main idea: 
        1. check total number of elements in arr if the same as r*c
        2. loop through by col , row of original arr and keep a temporary array to cache target col size
        3. every time temp size reached size of target col, append to answer
        4. return answer 
        """
        ro = len(nums)
        co = len(nums[0])
        if co* ro != r*c :
            return nums
        res = [[None] * c for _ in range(r)]
        i =0
        while i < r*c:
            v = nums[i//co][i%co] 
            temp.append(v)
            #print(temp)
            if len(temp) == c :
                res.append(temp)
                temp = []
            i += 1 
        return res 

space optimization

class Solution:
    def matrixReshape(self, nums: List[List[int]], r: int, c: int) -> List[List[int]]:
        """
        tc O(R*C) sp O(1)
        main idea: 
        1. check total number of elements in arr if the same as r*c
        2. loop through by col , row of original arr and keep a temporary array to cache target col size
        3. every time temp size reached size of target col, append to answer
        4. return answer 
        """
        ro = len(nums)
        co = len(nums[0])
        if co* ro != r*c :
            return nums
        res = [[None] * c for _ in range(r)]
        cnt = 0
        for i in range(ro):
            for j in range(co):
                res[cnt//c][cnt%c] = nums[i][j]
                cnt += 1
        return res