在python中跨一个轴查找3D数组中的最小值，并将另一个3D数组中的非对应值替换为0（不带循环）

假设我们有两个3D数组，A（x，y，z）和B（x，y，z），x，y，z是维度。我想识别A数组中z轴上的所有最小值，然后根据这些值及其索引选择B中的相应值，保留它们，并将其他值替换为零。

您可以稍微改变一下想法。在

A

中查找极小值的位置非常简单：

ind = np.expand_dims(np.argmin(A, axis=2), axis=2)

您可以执行以下操作之一：

最简单：创建

B

的替换项并填充相关元素：

 C = np.zeros_like(B)
 np.put_along_axis(C, ind, np.take_along_axis(B, ind, 2), 2)

同样的事情，但在适当的地方：

 values = np.take_along_axis(B, ind, 2)
 B[:] = 0
 np.put_along_axis(B, ind, values, 2)

将索引转换为掩码：

 mask = np.ones(B.shape, dtype=bool)
 np.put_along_axis(mask, ind, False, 2)
 B[mask] = 0

您可以使用合适的索引表达式来替换对

沿_轴取

和

沿_轴放置

的调用。特别是：

indx, indy = np.indices(A.shape[:-1])
indz = np.argmin(A, axis=-1)

然后，上面的示例转换为

新阵列：

 C = np.zeros_like(B)
 C[indx, indy, indz] = B[indx, indy, indz]

到位：

 values = B[indx, indy, indz]
 B[:] = 0
 B[indx, indy, indz] = values

蒙面的：

 mask = np.ones(B.shape, dtype=bool)
 mask[indx, indy, indz] = False
 B[mask] = 0