Python numpy ndarray 上的逐元素条件

Question

我有两个相同形状的 numpy ndarrays。

A = [[12, 25, 6],
    [28, 52, 74]]
B = [[100, 2, 4],
    [2, 12, 14]]

我的目标是将 B 中的值 <= 5 的每个元素替换为 A 中的 0。所以我的结果应该是：

# So C[0][0] = 12 because A[0][0] = 12 and B[0][0] >= 5
C = [[12, 0, 0],
    [0, 52, 74]]

有没有有效的方法来做到这一点？对于上下文，这是尝试对图像做一些背景减法，并将所有背景替换为黑色。

Answer 1

给你：

A = np.array([[12, 25, 6],[28, 52, 74]])
B = np.array([[100, 2, 4],[2, 12, 14]])

A = np.where(B <= 5, 0, A)

输出：

array([[12,  0,  0],
       [ 0, 52, 74]])

Answer 2

如果你想要一个新数组，我会这样做：

C = A.copy()
C[B <= 5] = 0

反正在我的机器上比 np.where() 快一点。

如果您不介意覆盖 A，只需 A[B <= 5] = 0。

Python element-wise condition on numpy ndarray