Splitting numpy multidimensional array based on indices stored in another array or list

Question

I have a numpy multidimensional array with shape = (12,2,3,3)

import numpy as np
arr = np.arange(12*2*3*3).reshape((12,2,3,3))

I need to select those elements based on the 2nd dimension where the dindices are stored in another list

indices = [0,1,0,0,1,1,0,1,1,0,1,1]

in one array, and the rest in another array. the output in either case should be in another array of shape (12,3,3)

arr2 = np.empty((arr.shape[0],*arr.shape[-2:]))

I could do it using a for loop

for i, ii in enumerate(indices):
    arr2[i] = arr[i, indices[ii],...]

However, I am searching for a one liner.

When I try indexing using the list as indices

test = arr[:,indices,...]

I get test of shape (12,12,3,3) instead of (12,3,3). Could you help me please?

My question is similar to this stackoverflow.com/questions/35497234/… but with indices instead of condition — Amr ALHOSSARY
– Amr ALHOSSARY, Commented Jul 10, 2021 at 3:35

jperezmartin · Accepted Answer · 2021-07-10 04:04:22Z

2

You can use np.arange for indexing the first dimension:

test = arr[np.arange(arr.shape[0]),indices,...]

or just the python range function:

test = arr[range(arr.shape[0]),indices,...]

answered Jul 10, 2021 at 3:48

jperezmartin

4174 silver badges20 bronze badges

Sign up to request clarification or add additional context in comments.

1 Answer 1