python - 简化 numpy argmin 的愚蠢循环

Question

如何简化我在此处以 numpy 样式编写的循环？

X     = np.random.random([10,15,20])
Y     = np.random.random([10,15,20,5])
Z     = np.zeros([10,15,5])

min_X = np.argmin(X,axis=2)
for i in range(10):
    for j in range(15):
        Z[i,j,:] = Y[i,j,min_X[i,j],:]

Answer 1

内置了 NumPy - np.take_along_axis为此(需要一些额外的步骤，因为它需要索引数组具有相同数量的暗淡)-

np.take_along_axis(Y,min_X[...,None,None],axis=2)[...,0,:]