Question

我想迭代一个numpy数组，只处理符合特定条件的元素。在下面的代码中，我想仅在元素大于1时才执行计算。

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

for i in range(0, a.shape[0]):
    for j in range(0, a.shape[1]):
        if a[i,j] > 1:
            a[i,j] = (a[i,j] - 3) * 5

是否可以使用单行代码而不是上面的双循环？也许让它更快？

Answer 1

方法＃1：使用布尔数组来索引：

>>> a = np.array([[1,3,5], [2,4,3], [1,2,0]])
>>> a[a > 1]  = (a[a > 1] - 3) * 5
>>> a
array([[ 1,  0, 10],
       [-5,  5,  0],
       [ 1, -5,  0]])

这会计算a > 1两次，但您可以将其分配给变量。（当然，在实践中，它不太可能成为瓶颈，但如果a足够大，那么内存可能是一个问题。）

方法＃2：使用np.where：

>>> a = np.array([[1,3,5], [2,4,3], [1,2,0]])
>>> np.where(a > 1, (a-3)*5, a)
array([[ 1,  0, 10],
       [-5,  5,  0],
       [ 1, -5,  0]])

这只会计算a > 1一次，但OTOH会为(ax-3)*5中的每个元素ax计算a，而不是只针对那些真正需要它的元素。< / p>

Answer 2

for index, x in np.ndenumerate(a):
    if x > 1:
        a[index] = (a[index] - 3) * 5

如果元素匹配条件，如何迭代Numpy数组并执行计算？

2 个答案: