如何使用NumPy掩码计算2D数组

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

现在的结果是一个一维数组，其形状为[3 6 5 3]。我该如何执行此操作并且不像之前纯Python示例中那样失去维度？

- gustavgans

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

- Divakar

---

谢谢！带上掩码应该是 a[mask] = a[mask] + 1 对吧？ - gustavgans

@gustavgans 这里的 a>0 是掩码（mask）。因此它等同于 a[mask] += 1。那个 += 是原地修改。另一种更明确的表达方式就像你所做的那样：a[mask] = a[mask]+ 1。 - Divakar