从 5D numpy 数组中选择相应的 3D 数组，其中包含第 4 维的索引

Posted 2023-02-14

【中文标题】从 5D numpy 数组中选择相应的 3D 数组，其中包含第 4 维的索引

【英文标题】：Selecting from 5D numpy array with a corresponding 3D array containing indices of the 4th dimension

【发布时间】：2022-01-22 00:36:28

【问题描述】：

我有一个包含值的 5D numpy 数组，并且想要获得一个维度更少的子数组，其中的值是根据包含第一个数组的第四维索引的 3D 数组选择的。例如，我有以下数组：

values = np.random.randn(3,4,5,10,2)
indices = np.random.randint(0,values.shape[3],size=values.shape[:3])

我找到了一种解决方案，但发现它相当复杂，并且更喜欢单线：

x = np.arange(values.shape[0])
y = np.arange(values.shape[1])
z = np.arange(values.shape[2])

result = values[x[:,None,None],y[None,:,None],z[None, None,:],indices,:]

有没有更好的办法来获取这个数组？

【参考方案1】：

您可以尝试以下方法：

indices = indices[..., None, None]
result = np.take_along_axis(values, indices, axis=3).squeeze(axis=3)

【讨论】：

谢谢！我尝试了很多take_along_axis，但没有弄清楚这两个Nons是必要的！

@Leon take_along_axis 要求两个数组具有相同的维数 - 因此indices 中的附加轴。