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

问题描述

我有 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,3))
# Cnew = np.where(A > B,D)-> ValueError: operands Could not be broadcast together with shapes (2,2) (2,3) (2,3) 

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

解决方法

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

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