查找沿另一个数组中的轴的最小值对应的数组

Question

我有两个三维数组,a和b,并且想要找到b的2D子阵列,其中a的元素沿着第三轴具有最小值,即

a=n.random.rand(20).reshape((5,2,2))
b=n.arange(20).reshape((5,2,2))
c=n.argmin(a,2) #indices with minimal value of a
d=n.zeros_like(c) #the array I want
for i in range(5):
  for j in range(2):
     d[i,j] = b[i,j,c[i,j]] 

有没有办法可以在没有双循环的情况下获得这些值？

我知道这个答案: 在numpy数组中将min值替换为另一个 但是如果我想让它适用于我的3D数组,我必须做很多重塑操作 - 我想知道是否有更简单的东西.

Answer 1

这是一种使用的方法fancy-indexing-

m,n,r = b.shape
d_out = b[np.arange(m)[:,None],np.arange(n),c]