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

时间:2012-09-27 14:14:53

标签: python numpy

假设我有以下Numpy数组:

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

我现在想要添加2,其中值不是0.最快的方法是什么?我必须操纵相当大的多维数组?

2 个答案:

答案 0 :(得分:5)

在我看来:

a[a!=0] += 2

应该有用。

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

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

当然,你可以自己保存掩码计算如果你可以在不同的地方重复使用相同的掩码(这是一个非常严格的约束):

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

当然,如果这足以成为瓶颈,您可以随时编写一个小的C / Fortran扩展来为您执行此操作(分别使用Cythonf2py)。这样可以避免掩码创建的开销。

答案 1 :(得分:1)

这取决于数组的大小,但您可以考虑使用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会使用所有可用核心。