根据不同数组的条件交换 2 个 numpy 数组

Question

我有 4 个数组，A、B、C、D。 A 和 B 的形状为 (n,n)，C/D 的形状为 (n,n,m)。我正在尝试设置它，以便当 A 的元素大于 B 时，长度为 m 的数组属于 C。本质上 C_new = np.where(A > B, C,D)，D_new = np.where(A < B , D, C)。然而，这给了我一个值错误 (operands could not be broadcast together with shapes)

我很好奇我是否可以在这里使用 where 而不是循环遍历每个元素？

编辑：示例：

A = np.ones((2,2))
B = 2*np.eye(2)
C = np.ones((2,2,3))
D = np.zeros((2,2,3))
# Cnew = np.where(A > B, C,D)-> ValueError: operands could not be broadcast together with shapes (2,2) (2,2,3) (2,2,3)

Cnew 在 (0,0) 和 (1,1) 索引中为零。

Answer 1

您需要在条件末尾添加一个新轴才能正确广播：

C_new = np.where((A > B)[..., np.newaxis], C, D)
D_new = np.where((A < B)[..., np.newaxis], D, C)

根据不同数组的条件交换 2 个 numpy 数组

Swap 2 numpy arrays based on condition from different arrays

python

numpy

array-broadcasting