Python 如何在numpy中对该循环进行矢量化？

我有这个阵列：

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

这些布尔指数：

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

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

我想找到布尔条件为true的数组中最大值的索引。因此，我：

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

它工作并返回1。但如果我有一个布尔索引数组：

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

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

---

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

对于

argmax

的特殊用例，您可以使用并将屏蔽值设置为负无穷大：

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

或者，您可以使用以下方式手动广播：

---

因此，您需要以下的矢量化版本：

In [302]: [np.ma.array(arr,mask=~c).argmax() for c in cond]
Out[302]: [1, 2]

cond

的实际尺寸是什么？如果行数与列数（或

arr

的长度）相比很小，那么这样的迭代可能并不昂贵

使用

tile

看起来不错。我在这里稍作修改：

In [308]: np.ma.array(np.tile(arr,(cond.shape[0],1)),mask=~cond).argmax(axis=1)
Out[308]: array([1, 2], dtype=int32)

正如预期的那样，列表理解时间随着

cond

的行进行缩放，而平铺方法只比单行情况慢一点。但是随着时间的推移，

92.7µs

这种掩蔽阵列方法比

arr.argmax（）

慢得多。屏蔽会增加很多开销

where

版本要快一点

np.where(cond, arr, -100).argmax(1)  # 20 µs

建议删除答案

(arr*cond).argmax(1)   # 8 µs

---

甚至更快。正如建议的那样，如果存在负的

arr

值，则它不起作用。但它可能可以调整以处理这些问题。

沿轴应用

在封面下进行迭代。所以我不希望有任何速度优势。