如何在 numpy 中对这个循环进行矢量化？答案

【问题标题】：How do I vectorize this loop in numpy?如何在 numpy 中对这个循环进行矢量化？
【发布时间】：2015-08-01 23:24:57
【问题描述】：

我有这个数组：

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

【问题讨论】：

标签： python numpy

【解决方案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])

【讨论】：

【解决方案2】：

所以你想要一个矢量化版本的：

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

cond 的实际尺寸是多少？如果行数比列数少（或arr 的长度），这样的迭代可能并不昂贵。

https://stackoverflow.com/a/31767220/901925 使用 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 值，它就不起作用。但它可能可以调整以处理这些问题。

【讨论】：

【解决方案3】：

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

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


def multi_slice_max(bool_arr , x ):
    return np.ma.array(x, mask=~bool_arr).argmax()

np.apply_along_axis(multi_slice_max , 1 , cond , arr)

【讨论】：

apply_along_axis 在幕后进行迭代。所以我不希望有任何速度优势。