Numpy:根据条件得到最小值的索引

时间:2014-03-05 15:56:08

标签: python numpy

我有一个数组:

array([[ 10, -1],
       [ 3,  1],
       [ 5, -1],
       [ 7,  1]])

我想要的是获取第一列中值最小的行的索引,第二列中的值为-1。

基本上,np.argmin()第二列的条件等于-1(或任何其他值)。

在我的示例中,我希望得到2,这是[ 5, -1]的索引。

我很确定这是一种简单的方法,但我找不到它。

3 个答案:

答案 0 :(得分:2)

import numpy as np

a = np.array([
    [10, -1],
    [ 3,  1],
    [ 5, -1],
    [ 7,  1]])

mask = (a[:, 1] == -1)

arg = np.argmin(a[mask][:, 0])
result = np.arange(a.shape[0])[mask][arg]
print result

答案 1 :(得分:2)

np.argwhere(a[:,1] == -1)[np.argmin(a[a[:, 1] == -1, 0])]

答案 2 :(得分:0)

这样效率不高但是如果你有一个相对较小的数组并想要一个单行解决方案:

>>> a = np.array([[ 10, -1],
...               [ 3,  1],
...               [ 5, -1],
...               [ 7,  1]])
>>> [i for i in np.argsort(a[:, 0]) if a[i, 1] == -1][0]
2