如何通过单个NumPy调用实现三维数组的指定索引向量提取

给定形状为(N, L, H)的三维数组data，以及长度为N的一维索引数组inds（每个元素取值范围为[0, L-1]），需要通过单次NumPy调用完成以下循环的逻辑：

import numpy as np

N, L, H = 5, 3, 2
data = np.arange(N * L * H).reshape(N, L, H)
inds = np.arange(N) % L
    
indexed_data = []
for x, i in zip(data, inds):
    indexed_data.append(x[i])
y = np.array(indexed_data)

assert y.shape == (N, H)

解决方案 #

方法1：使用np.take #

利用np.take指定axis=1，对每个N维度的子数组沿L维度取对应索引的元素，之后挤压掉多余的维度：

y = np.take(data, inds, axis=1).squeeze(axis=1)
assert y.shape == (N, H)

方法2：花式索引（更简洁高效） #

构造匹配的N维度索引，结合inds直接定位目标元素，无需额外维度处理：

n_indices = np.arange(N)
y = data[n_indices, inds]
assert y.shape == (N, H)

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