Reshaping 3D Numpy Array to a 2D array

Question

I have the following 3D array in Numpy:

a = np.array([[[1,2],[3,4]], [[5,6],[7,8]], [[9,10],[11,12]],[[13,14],[15,16]]])

when I write

b = np.reshape(a, [4,4])

The 2D resulting array will look like

 [[ 1  2  3  4]
  [ 5  6  7  8]
  [ 9 10 11 12]
  [13 14 15 16]]

However, I want it to be in this shape:

 [[ 1  2  5  6]
  [ 3  4  7  8]
  [ 9 10 13 14]
  [11 12 15 16]]

How can I do this efficiently in Python/Numpy?

Answer 1

Another approach using just np.hstack and np.vstack:

In [98]: a
Out[98]: 
array([[[ 1,  2],
        [ 3,  4]],

       [[ 5,  6],
        [ 7,  8]],

       [[ 9, 10],
        [11, 12]],

       [[13, 14],
        [15, 16]]])

In [99]: s0, s1, s2, s3 = range(a.shape[0])

In [100]: np.vstack((np.hstack((a[s0], a[s1])), np.hstack((a[s2], a[s3]))))
Out[100]: 
array([[ 1,  2,  5,  6],
       [ 3,  4,  7,  8],
       [ 9, 10, 13, 14],
       [11, 12, 15, 16]])

Realizing the fact that your aim is to squash the first two slices of your original array into one slice and the next two into another slice and so on.

And you could also just replace the np.vstack and np.hstack with their fastest cousin np.concatenate, if you're concerned about performance.

P.S.: This approach creates a new array leaving your original one unchanged.

Answer 2

Reshape to split the first axis into two, permute axes and one more reshape -

a.reshape(2,2,2,2).transpose(0,2,1,3).reshape(4,4)
a.reshape(2,2,2,2).swapaxes(1,2).reshape(4,4)

Making it generic, would become -

m,n,r = a.shape
out = a.reshape(m//2,2,n,r).swapaxes(1,2).reshape(-1,2*r)

Sample run -

In [20]: a
Out[20]: 
array([[[ 1,  2],
        [ 3,  4]],

       [[ 5,  6],
        [ 7,  8]],

       [[ 9, 10],
        [11, 12]],

       [[13, 14],
        [15, 16]]])

In [21]: a.reshape(2,2,2,2).swapaxes(1,2).reshape(4,4)
Out[21]: 
array([[ 1,  2,  5,  6],
       [ 3,  4,  7,  8],
       [ 9, 10, 13, 14],
       [11, 12, 15, 16]])