如何在numpy中将此循环向量化？

Question

我有这个数组:

arr = np.array([3, 7, 4])

这些布尔索引:

cond = np.array([False, True, True])

我想在布尔条件为真的数组中找到最大值的索引.所以我这样做:

np.ma.array(arr, mask=~cond).argmax()

哪个工作并返回1.但如果我有一个布尔索引数组:

cond = np.array([[False, True, True], [True, False, True]])

是否有一个矢量化/ numpy方式迭代布尔索引数组返回[1,2]？

Answer 1

对于 的特殊用例argmax，您可以使用np.where掩码值并将其设置为负无穷大：

>>> inf = np.iinfo('i8').max
>>> np.where(cond, arr, -inf).argmax(axis=1)
array([1, 2])

或者，您可以使用以下命令手动广播np.tile：

>>> np.ma.array(np.tile(arr, 2).reshape(2, 3), mask=~cond).argmax(axis=1)
array([1, 2])