如何用索引数组对numpy多维数组最后一维做通用索引

方案1：使用np.take_along_axis（最简洁，推荐） #

这是Numpy专门为这类按轴匹配索引的场景设计的API，适配任意维度n，代码如下：

import numpy as np
F = np.take_along_axis(A, I[..., np.newaxis], axis=-1).squeeze(axis=-1)

• 逻辑说明：I[..., np.newaxis]会给索引数组I增加最后一个长度为1的维度，让它的维度数和A保持一致，匹配take_along_axis的参数要求；axis=-1指定索引作用在最后一维；最后用squeeze去掉多余的长度为1的维度，得到和I形状完全一致的结果数组F。
• 版本要求：Numpy 1.15.0及以上版本支持该API。

方案2：使用高级索引构造（兼容旧版本Numpy） #

如果使用的Numpy版本低于1.15，可以通过构造前n维的索引网格实现：

import numpy as np
# 生成前n个维度的索引网格
pre_indices = np.indices(I.shape)
# 拼接最后一维的索引I，传入A做高级索引
F = A[tuple(pre_indices) + (I,)]

• 逻辑说明：np.indices(I.shape)会生成和I前n维一一对应的索引数组，拼接上I作为最后一维的索引后，即可实现每个位置的前n维索引和I中的最后一维索引一一匹配，直接得到目标形状的F。

示例验证 #

以n=3的场景为例验证效果：

# 构造测试数据，A形状为(2,3,4)，对应a1=2,a2=3,b=4
A = np.arange(2*3*4).reshape(2,3,4)
# 构造索引数组I，形状为(2,3)，元素范围0~3
I = np.random.randint(0,4, size=(2,3))
# 两种方案计算F
F1 = np.take_along_axis(A, I[..., np.newaxis], axis=-1).squeeze(-1)
F2 = A[tuple(np.indices(I.shape)) + (I,)]
# 验证结果一致
print(np.array_equal(F1, F2)) # 输出True

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