如何使用 3d 数组和 2d 数组 numpy 进行掩码

如何使用 1d 数组从 3d 数组中选择一组元素。

#These are my 3 data types
# A = numpy.ndarray[numpy.ndarray[float]]
# B1 = numpy.ndarray[numpy.ndarray[numpy.ndarray[float]]]
#B2=numpy.ndarray[numpy.ndarray[numpy.ndarray[float]]]
#I want to choose values from A based on values from B1 in the B2

这是我尝试过的，但它返回了所有False:

A2[i]=image_values[updated_image_values==initial_means[i]]

Example:

A=[[1,1,1][2,2,2]]
B=[[[1,1,1],[2,3,4]],[[2,2,2],[1,1,1]],[[1,1,1],[2,2,2]]]
B2=[[[2,2,2],[9,3,21]],[[22,0,-2],[-1,-1,1]],[[1,-1,-1],[10,0,2]]]

#A2 is calculated as the means of the B2 values that correspond 
#to it's value according to B

所以，要计算A2我们使用检查 B2 中的值等于 A 中的值。因此，对于第一个索引A[0], B[0][0],B[1][1] and B[2][0]等于A[0]。因此对于A2[0]，我们得到相应的值B in B2并使用它们来计算每个指数的平均值：

#A2[0][0]=(B2[0][0][0]+B2[1][1][0]+B2[2][0][0]) /3 = 0.67

#A2[1][2]=(B2[1][0][2]+B2[2][1][2]) /2 = 0

#After doing this for every A2 value, A2 should be:

A2=[[0.67,0,0.67],[16,0,0]]

这是一种矢量化方法np.add.reduceat -

idx = np.argwhere((B == A[:,None,None]).all(-1))
B2_indexed = B2[idx[:,1],idx[:,2]]
_,start, count = np.unique(idx[:,0],return_index=1,return_counts=1)
out = np.add.reduceat(B2_indexed,start)/count.astype(float)[:,None]

或者，我们可以通过避免创建来节省一点内存4D面具与一个3D掩码代替获取idx，就像这样-

dims = np.maximum(B.max(axis=(0,1)),A.max(0))+1
A_reduced = np.ravel_multi_index(A.T,dims)
B_reduced = np.ravel_multi_index(B.T,dims)
idx = np.argwhere(B_reduced.T == A_reduced[:,None,None])

这是另一种单循环方法 -

out = np.empty(A.shape)
for i in range(A.shape[0]):
    r,c = np.where((B == A[i]).all(-1))    
    out[i] = B2[r,c].mean(0)