CODEWITHSUNDEEP
numpyindexingnumpy-ndarray
ma0ho
ma0ho

Reputation: 633

Numpy: Select by index array along an axis

I'd like to select elements from an array along a specific axis given an index array. For example, given the arrays

a = np.arange(30).reshape(5,2,3)
idx = np.array([0,1,1,0,0])

I'd like to select from the second dimension of a according to idx, such that the resulting array is of shape (5,3). Can anyone help me with that?

Upvotes: 5

Views: 3350

Answers (2)

Kevin
Kevin

Reputation: 3348

You could use fancy indexing

a[np.arange(5),idx]

Output:

array([[ 0,  1,  2],
       [ 9, 10, 11],
       [15, 16, 17],
       [18, 19, 20],
       [24, 25, 26]])

To make this more verbose this is the same as:

x,y,z = np.arange(a.shape[0]), idx, slice(None)
a[x,y,z]

x and y are being broadcasted to the shape (5,5). z could be used to select any columns in the output.

Upvotes: 3

sjw
sjw

Reputation: 6543

I think this gives the results you are after - it uses np.take_along_axis, but first you need to reshape your idx array so that it is also a 3d array:

a = np.arange(30).reshape(5, 2, 3)
idx = np.array([0, 1, 1, 0, 0]).reshape(5, 1, 1)
results = np.take_along_axis(a, idx, 1).reshape(5, 3)

Giving:

[[ 0  1  2]
 [ 9 10 11]
 [15 16 17]
 [18 19 20]
 [24 25 26]]

Upvotes: 4

Related Questions