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）。这将避免创建掩码的开销。

Answer 2

这取决于数组的大小，但您可以考虑为此使用numexpr：

# from your exemmple
>>> a = array([[3, 5, 0], [7, 0, 2]])
>>> %timeit a[a!=0] += 2
100000 loops, best of 3: 18.6 us per loop
>>> timeit numexpr.evaluate("a + 2 * (a != 0)")
10000 loops, best of 3: 42.6 us per loop

但是使用更大的数组：

# make a big array with 10% of zeros :
a = np.random.rand(10000)
a[a<0.1] = 0
# same test:
>>> timeit a[a!=0] += 2
1000 loops, best of 3: 364 us per loop
>>> timeit numexpr.evaluate("a + 2 * (a != 0)")
10000 loops, best of 3: 119 us per loop

这是一个单核。如果可能，Numepxr 会使用所有可用的内核。