如何使用numpy掩码计算二维数组

Question

我有一个二维数组，如果值大于 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 示例那样丢失尺寸？

Answer 1

如果 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]])

如何使用numpy掩码计算二维数组

how to calculate a 2D array with numpy mask

python

arrays

numpy

vectorization