python - bool 索引与 np.where

Question

好吧，假设我有一个 numpy 数组 arr 和一个相同形状的 bool 数组 mask(例如 mask = arr >= 20)

我想要一个包含 arr 的所有值的数组，其中 mask 为 True。我真的不关心顺序(我会在之后计算这个总和)

根据我从 the numpy doc 收集到的信息，我可以只使用 bool 索引:

arr[mask]

尽管如此，在互联网上，我看到了很多这样的代码:

arr[np.where(mask)]

我认为，它的作用相同，但使用索引数组。

这两行真的做同样的事情吗？如果是这样，其中一个更快吗？

Answer 1

至于性能:为什么不简单地衡量它呢？举个简单的例子:

In [11]: y = np.arange(35).reshape(5,7)

In [12]: mask = (y % 2 == 0)

In [13]: mask
Out[13]:
array([[ True, False,  True, False,  True, False,  True],
       [False,  True, False,  True, False,  True, False],
       [ True, False,  True, False,  True, False,  True],
       [False,  True, False,  True, False,  True, False],
       [ True, False,  True, False,  True, False,  True]])

然后 %timeit:

In [14]: %timeit y[mask]
534 ns ± 1.61 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [15]: %timeit y[np.where(mask)]
2.18 µs ± 16.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

毫不奇怪 - 即使两条线之间没有功能差异 - 函数调用开销也会使 np.where 变慢。至于“它们是否相同”？不完全是。来自 np.where 文档字符串:

where(condition, [x, y]):Return elements chosen from x or y depending on condition.

Note: When only condition is provided, this function is a shorthand fornp.asarray(condition).nonzero(). Using nonzero directly should bepreferred, as it behaves correctly for subclasses. The rest of thisdocumentation covers only the case where all three arguments areprovided.

回顾一下这个例子:

虽然 y[mask] 直接 选择 y 的所有匹配(True)元素， np.where(mask) 绕道计算掩码中 True 元素的所有(此处为 2D)索引位置:

In [26]: np.where(mask)
Out[26]:
(array([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4, 4], dtype=int64),
 array([0, 2, 4, 6, 1, 3, 5, 0, 2, 4, 6, 1, 3, 5, 0, 2, 4, 6], dtype=int64))

换句话说:直接使用 bool 掩码不仅更简单，而且避免了很多额外的计算。