未知前置维度时，如何用argmin索引numpy数组X与Y？

核心思路 #

使用Numpy的np.take_along_axis()函数，它专门用于沿指定轴根据索引数组提取元素，完全适配任意数量的前置维度，无需手动处理固定维度的索引逻辑。

1. 提取Y中对应argmin的元素 #

Y的形状为(..., 12)，idx是np.argmin(Y, axis=-1)得到的形状为(...)的索引数组：

import numpy as np

# 假设已有Y和idx
Y_min = np.take_along_axis(Y, idx[..., np.newaxis], axis=-1).squeeze(-1)

• idx[..., np.newaxis]：给索引数组扩展一个维度，使其与Y的最后一维（长度12）匹配，形状变为(..., 1)
• take_along_axis：沿axis=-1（Y的最后一维）提取对应索引的元素，结果形状为(..., 1)
• squeeze(-1)：去掉最后一个多余维度，最终结果形状为(...)

2. 提取X中对应argmin的元素 #

X的形状为(..., 12, 3)，需要沿倒数第二维（长度12的维度，对应axis=-2）提取：

# 假设已有X和idx
X_selected = np.take_along_axis(X, idx[..., np.newaxis, np.newaxis], axis=-2).squeeze(-2)

• idx[..., np.newaxis, np.newaxis]：将索引数组扩展为(..., 1, 1)，与X的倒数第二维匹配，同时兼容最后一维（长度3）的广播
• take_along_axis：沿axis=-2提取元素，结果形状为(..., 1, 3)
• squeeze(-2)：去掉倒数第二维的多余维度，最终结果形状为(..., 3)

替代简化写法（利用广播） #

对于X的提取，也可以用更简洁的维度扩展方式：

X_selected = np.take_along_axis(X, np.expand_dims(idx, (-2, -1)), axis=-2).squeeze(-2)

效果和上述代码完全一致，仅用np.expand_dims一次性扩展两个维度，可读性更强。

内容的提问来源于stack exchange，提问作者John