python - 我如何在 numpy 中向量化这个循环？

Question

我有这个数组:

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

还有这些 bool 索引:

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

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

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

有效并返回 1。但是如果我有一个 bool 索引数组:

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

是否有一种向量化/numpy 的方法来遍历 bool 索引数组以返回 [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])