将numpy索引应用于矩阵

Question

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 😊

Answer 1

是因为argmin返回每一行的列索引（其中axis=1 ），因此您需要访问这些特定列的每一行：

a[range(a.shape[0]), amin_index]

Answer 2

One way is to pass in the array of row indexes (eg [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.

Answer 3

Why not simply do np.amin(np.abs(a), axis=1) , it's much simpler you don't need the intermediate amin_index array via argmin?

Numpy's reference page is an excellent resource, see "Indexing".

Timing is always useful:

In [3]: a=np.random.rand(4000, 4000)-.5

In [4]: %timeit np.amin(np.abs(a), axis=1)
10 loops, best of 3: 128 ms per loop

In [5]: %timeit a[np.arange(a.shape[0]), np.argmin(np.abs(a), axis=1)]
10 loops, best of 3: 135 ms per loop