One using iterator.

```
class Solution(object):
def matrixReshape(self, nums, r, c):
"""
:type nums: List[List[int]]
:type r: int
:type c: int
:rtype: List[List[int]]
"""
A = (x for row in nums for x in row)
return [[next(A) for i in range(c)] for j in range(r)] \
if c * r == len(nums) * len(nums[0]) else nums
```

And the other using list to store all `num`

in `nums`

.

```
class Solution(object):
def matrixReshape(self, nums, r, c):
"""
:type nums: List[List[int]]
:type r: int
:type c: int
:rtype: List[List[int]]
"""
A = [x for row in nums for x in row]
return [A[i*c:(i+1)*c] for i in range(r)] if c * r == len(A) else nums
```