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?
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]]
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.
a[np.r_[:5], idx].