在一个轴上的3D数组中查找最小值，并在python中用0替换另一个3D数组中的不对应值而没有循环

问题描述

比方说，我们有两个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)

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

1. 最简单：创建B的替换并填充相关元素：

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

2. 相同，但就地：

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

3. 将索引转换为掩码：

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

您可以使用合适的索引表达式替换对take_along_axis和put_along_axis的调用。特别是：

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

然后将以上示例转换为

1. 新数组：

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

2. 就地：

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

3. 蒙版：

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