我怎么能在某些条件下获得 numpy 数组索引

Question

我遇到这样的问题： 假设我有这样的数组： a = np.array([[1,2,3,4,5,4,3,2,1],]) label = np.array([[1,0,1,0,0,1,1,0,1],]) 我需要获取a 的索引，在该位置label 的元素值为1，a 的值是导致label 为1 的最大数量。

在上面的例子中，label 为 1 的索引可能令人困惑：0、2、5、6、8，它们对应的a 的值因此为：1、3、4、3， 1，其中4是最大的，因此我需要得到5的结果，即a中数字4的索引。我怎么能用 numpy 做到这一点？

Answer 1

获取1s 索引，比如idx，然后用它索引a，获取max 索引，最后通过索引idx 将其追溯到原始顺序-

idx = np.flatnonzero(label==1)
out = idx[a[idx].argmax()]

示例运行 -

# Assuming inputs to be 1D
In [18]: a
Out[18]: array([1, 2, 3, 4, 5, 4, 3, 2, 1])

In [19]: label
Out[19]: array([1, 0, 1, 0, 0, 1, 1, 0, 1])

In [20]: idx = np.flatnonzero(label==1)

In [21]: idx[a[idx].argmax()]
Out[21]: 5

对于作为整数的a 和作为0s 和1s 的数组的label，我们可以进一步优化，因为我们可以根据其中的值范围缩放a，就像这样-

(label*(a.max()-a.min()+1) + a).argmax()

此外，如果a 只有正数，它会简化为 -

(label*(a.max()+1) + a).argmax()

正整数的时序较大a -

In [115]: np.random.seed(0)
     ...: a = np.random.randint(0,10,(100000))
     ...: label = np.random.randint(0,2,(100000))

In [117]: %%timeit
     ...: idx = np.flatnonzero(label==1)
     ...: out = idx[a[idx].argmax()]
1000 loops, best of 3: 592 µs per loop

In [116]: %timeit (label*(a.max()-a.min()+1) + a).argmax()
1000 loops, best of 3: 357 µs per loop

# @coldspeed's soln
In [120]: %timeit np.ma.masked_where(~label.astype(bool), a).argmax()
1000 loops, best of 3: 1.63 ms per loop

# won't work with negative numbers in a
In [119]: %timeit (label*(a.max()+1) + a).argmax()
1000 loops, best of 3: 292 µs per loop

# @klim's soln (won't work with negative numbers in a)
In [121]: %timeit np.argmax(a * (label == 1))
1000 loops, best of 3: 229 µs per loop

Answer 2

这是最简单的方法之一。

>>> np.argmax(a * (label == 1))
5
>>> np.argmax(a * (label == 1), axis=1)
array([5])

Coldspeed 的方法可能需要更多时间。

Answer 3

您可以使用掩码数组：

>>> np.ma.masked_where(~label.astype(bool), a).argmax()
5