python - 根据来自不同数组的条件交换 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)