CODEWITHSUNDEEP
pythonarraysnumpysub-array
StatsNoob
StatsNoob

Reputation: 390

How to rearrange an array by subarray elegantly in numpy?

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!!

Upvotes: 2

Views: 818

Answers (2)

Alleo
Alleo

Reputation: 8518

Using einops:

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

And I would recommend to give meaningful names to axes (instead of x y z) depending on the context (e.g. 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]])

Upvotes: 4

Divakar
Divakar

Reputation: 221534

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]])

Upvotes: 3

Related Questions