Curious Curious - 4 months ago 19
Python Question

How to flatten only some dimensions of a numpy array

Is there a quick way to "sub-flatten" or flatten only some of the first dimensions in a numpy array?

For example, given a numpy array of dimensions

(50,100,25)
, the resultant dimensions would be
(5000,25)


Thanks

Answer

Take a look at numpy.reshape .

>>> arr = numpy.zeros((50,100,25))
>>> arr.shape
# (50, 100, 25)

>>> new_arr = arr.reshape(5000,25)
>>> new_arr.shape   
# (5000, 25)

# One shape dimension can be -1. 
# In this case, the value is inferred from 
# the length of the array and remaining dimensions.
another_arr = arr.reshape(-1, arr.shape[-1])
another_arr.shape
# (5000, 25)