0

我想迭代一个 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

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

4

2 回答 2

3

方法#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而不是只对那些真正需要它的元素进行计算。

于 2013-05-07T03:59:13.390 回答
0 
for index, x in np.ndenumerate(a):
    if x > 1:
        a[index] = (a[index] - 3) * 5
于 2013-05-07T03:53:00.353 回答