如何在numpy中通过子数组优雅地重新排列数组？

Question

Let's say I have a 3-D array:

[[[0,1,2],
  [0,1,2],
  [0,1,2]],

 [[3,4,5],
  [3,4,5],
  [3,4,5]]]

And I want to rearrange this by the columns:

[[0,1,2,3,4,5],
 [0,1,2,3,4,5],
 [0,1,2,3,4,5]]

What would be an elegant python numpy code for doing this for essentially a 3-D np.array of arbitrary shape and depth? Could there be a fast method that bypasses for loop? All the approaches I made were terribly adhoc and brute they were basically too slow and useless...

Thanks!!

Answer 1

Swap axes and reshape -

a.swapaxes(0,1).reshape(a.shape[1],-1)

Sample run -

In [115]: a
Out[115]: 
array([[[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]],

       [[3, 4, 5],
        [3, 4, 5],
        [3, 4, 5]]])

In [116]: a.swapaxes(0,1).reshape(a.shape[1],-1)
Out[116]: 
array([[0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5]])

Answer 2

Using einops:

einops.rearrange(a, 'x y z -> y (x z) ')

And I would recommend to give meaningful names to axes (instead of xyz) depending on the context (eg time, height, etc.). This will make it easy to understand what the code does

In : einops.rearrange(a, 'x y z -> y (x z) ')
Out:
array([[0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5]])