HansSnah HansSnah - 2 months ago 6
Python Question

Apply numpy index to matrix

I have spent the last hour trying to figure this out

Suppose we have

import numpy as np
a = np.random.rand(5, 20) - 0.5
amin_index = np.argmin(np.abs(a), axis=1)
> [ 0 12 5 18 1] # or something similar

this does not work:


So, in essence, I need to find the minima along a certain axis for the array np.abs(a), but then extract the values from the array a at these positions. How can I apply an index to just one axis?

Probably very simple, but I can't get it figured out. Also, I can't use any loops since I have to do this for arrays with several million entries.


One way is to pass in the array of row indexes (e.g. [0,1,2,3,4]) and the list of column indexes for the minimum in each corresponding row (your list amin_index).

This returns an array containing the value at [i, amin_index[i]] for each row i:

>>> a[np.arange(a.shape[0]), amin_index]
array([-0.0069325 ,  0.04268358, -0.00128002, -0.01185333, -0.00389487])

This is basic indexing (rather than advanced indexing), so the returned array is actually a view of a rather than a new array in memory.