如何使用numpy mask计算2D数组

时间:2015-06-14 08:23:30

标签: python arrays numpy vectorization

我有一个2维数组,如果值大于0,我想做一个操作(例如x + 1)。 在普通的python中是这样的:

a = [[2,5], [4,0], [0,2]]
for x in range(3):
    for y in range(2):
        if a[x][y] > 0:
            a[x][y] = a[x][y] + 1 

a的结果是 [[3,6],[5,0],[0,3]] 。这就是我想要的。

现在我想阻止嵌套循环并尝试使用numpy这样的东西:

a = np.array([[2,5], [4,0], [0,2]])
mask = (a > 0)
a[mask] + 1

现在结果为1维,阵列的形状为 [3 6 5 3] 。我怎么能做这个操作,不要像之前的普通python例子那样松开维度?

1 个答案:

答案 0 :(得分:2)

如果a是一个numpy数组,你可以简单地做 -

a[a>0] +=1

示例运行 -

In [335]: a = np.array([[2,5], [4,0], [0,2]])

In [336]: a
Out[336]: 
array([[2, 5],
       [4, 0],
       [0, 2]])

In [337]: a[a>0] +=1

In [338]: a
Out[338]: 
array([[3, 6],
       [5, 0],
       [0, 3]])