python - 选择性 Numpy 数组操作(取决于值)

Question

假设我有以下 Numpy 数组:

array([[3, 5, 0], [7, 0, 2]])

我现在想在值不为 0 的地方加 2。最快的方法是什么？我必须操纵相当大的多维数组？

Answer 1

在我看来:

a[a!=0] += 2

应该可以。

(对于非零性测试的有限情况)，您可能能够加快速度(您需要 timeit 才能看到):

mask = a.astype(bool)
a[mask] += 2

---

当然，您可以节省掩码计算如果您可以在不同的地方重复使用相同的掩码(这是一个非常严格的约束):

mask = a != 0
a[mask] += 2
#some more code ...
a[mask] *= 3
#more code ...

---

当然，如果这足以成为瓶颈，您总是可以编写一些 C/Fortran 扩展来为您做这件事(分别使用 Cython 或 f2py) .这将避免掩码创建的开销。