在numpy数组中将特定值设置为零

Question

我有一个形状为(4096,4096) 的numpy 数组/矩阵和一个应该设置为零的元素数组。我发现函数numpy.in1d 工作正常，但对我的计算来说非常慢。我想知道是否存在一些更快的方法来执行此操作，因为我需要在大量矩阵上重复此操作，因此每次优化都有帮助。

示例如下：

numpy 数组如下所示：

npArr = np.array([
    [1, 4, 5, 5, 3],
    [2, 5, 6, 6, 1],
    [0, 0, 1, 0, 0],
    [3, 3, 2, 4, 3]])

另一个数组是：

arr = np.array([3,5,8])

替换后的 numpy 数组 npArr 应该是这样的：

array([[ 1,  4,  0,  0,  0],
       [ 2,  0,  6,  6,  1],
       [ 0,  0,  1,  0,  0],
       [ 0,  0,  2,  4,  0]])

Answer 1

如果您有numba，您可以使用不需要中间掩码的自定义函数来解决此问题：

import numpy as np
import numba as nb

@nb.njit
def replace_where(arr, needle, replace):
    arr = arr.ravel()
    needles = set(needle)
    for idx in range(arr.size):
        if arr[idx] in needles:
            arr[idx] = replace

这为您的示例提供了正确的结果：

npArr = np.array([[1, 4, 5, 5, 3],
                  [2, 5, 6, 6, 1],
                  [0, 0, 1, 0, 0],
                  [3, 3, 2, 4, 3]])

arr = np.array([3,5,8])

replace_where(npArr, arr, 0)
print(npArr)
# array([[1, 4, 0, 0, 0],
#        [2, 0, 6, 6, 1],
#        [0, 0, 1, 0, 0],
#        [0, 0, 2, 4, 0]])

而且它应该非常非常快。我对几个数组大小进行了计时，它比np.in1d 快5-20 倍（取决于大小，尤其是arr 大小）。

Answer 2

这是使用np.searchsorted 的替代方法-

def in1d_alternative_2D(npArr, arr):
    idx = np.searchsorted(arr, npArr.ravel())
    idx[idx==len(arr)] = 0
    return arr[idx].reshape(npArr.shape) == npArr

它假定arr 被排序。如果不是，我们需要排序，然后使用posted方法。

示例运行 -

In [90]: npArr = np.array([[1, 4, 5, 5, 3],
    ...:     [2, 5, 6, 6, 1],
    ...:     [0, 0, 1, 0, 0],
    ...:     [3, 3, 2, 14, 3]])
    ...: 
    ...: arr = np.array([3,5,8])
    ...: 

In [91]: in1d_alternative_2D(npArr, arr)
Out[91]: 
array([[False, False,  True,  True,  True],
       [False,  True, False, False, False],
       [False, False, False, False, False],
       [ True,  True, False, False,  True]], dtype=bool)

In [92]: npArr[in1d_alternative_2D(npArr, arr)] = 0

In [93]: npArr
Out[93]: 
array([[ 1,  4,  0,  0,  0],
       [ 2,  0,  6,  6,  1],
       [ 0,  0,  1,  0,  0],
       [ 0,  0,  2, 14,  0]])

针对numpy.in1d进行基准测试

使用np.in1d 的等效解决方案是：

np.in1d(npArr, arr).reshape(npArr.shape)

让我们对我们提出的建议进行计时，并验证问题中提到的尺寸的结果。

In [85]: # (4096, 4096) shaped 'npArr' and search array 'arr' of 1000 elems
    ...: npArr = np.random.randint(0,10000,(4096,4096))
    ...: arr = np.sort(np.random.choice(10000, 1000, replace=0 ))
    ...: 

In [86]: out1 = np.in1d(npArr, arr).reshape(npArr.shape)
    ...: out2 = in1d_alternative_2D(npArr, arr)
    ...: 

In [87]: np.allclose(out1, out2)
Out[87]: True

In [88]: %timeit np.in1d(npArr, arr).reshape(npArr.shape)
1 loops, best of 3: 3.04 s per loop

In [89]: %timeit in1d_alternative_2D(npArr, arr)
1 loops, best of 3: 1 s per loop

Answer 3

使用 numpy 广播的另一种解决方案：

np.min(np.where(npArr[None,:,:] == arr[:,None,None], 0, a),0)
Out[730]: 
array([[1, 4, 0, 0, 0],
       [2, 0, 6, 6, 1],
       [0, 0, 1, 0, 0],
       [0, 0, 2, 4, 0]])