HansSnah - 1 year ago 138
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)
print(amin_index)
> [ 0 12  5 18  1] # or something similar
``````

this does not work:

``````a[amin_index]
``````

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.
thanks

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