Numpy：根据条件获取最小值的索引

Question

我有一个如下的数组：

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

我想要的是获取第一列最小值且第二列为-1（或其他值）时行的索引。

基本上，使用np.argmin()并加上第二列等于-1的条件（或其他条件）即可。

在我的例子中，我想要获取[ 5, -1]的索引2。

我相信有一种简单的方法，但我找不到它。

- cgf

---

你可以像这样预过滤数组：np.argmin([i for i in a if i[1] == -1])，但它仍然会在数组[[10，-1]，[5，-1]]上返回一个。 - njzk2

是的，我已经尝试过这些方法了。不过我需要的是原始数组的索引。 - cgf

Answer 1

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

- YXD

Answer 2

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

- Blaz Bratanic

---

不确定问题出在哪里，但如果我尝试使用np.argwhere(d[:,1] == 1)[np.argmin(d[d[:, 1] == 1])]（从-1改为1），我会得到错误的结果。 - cgf

我认为你需要选择适当的列，例如像 np.argwhere(a[:,1] == -1)[np.argmin(a[a[:, 1] == -1,0])] 这样（未经测试）。 - DSM

是的，我在argmin中错过了列选择。谢谢。 - Blaz Bratanic

Answer 3

这并不高效，但如果您有一个相对较小的数组并且想要一行解决方案：

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

- bogatron