Numpy：沿轴按索引数组选择

【问题标题】：Numpy: Select by index array along an axisNumpy：沿轴按索引数组选择
【发布时间】：2021-07-20 16:29:13
【问题描述】：

我想在给定索引数组的情况下沿特定轴从数组中选择元素。例如，给定数组

a = np.arange(30).reshape(5,2,3)
idx = np.array([0,1,1,0,0])

我想根据idx从a的第二维中进行选择，这样得到的数组的形状是(5,3)。谁能帮我解决这个问题？

【问题讨论】：

【解决方案1】：

你可以使用花哨的索引

a[np.arange(5),idx]

输出：

array([[ 0,  1,  2],
       [ 9, 10, 11],
       [15, 16, 17],
       [18, 19, 20],
       [24, 25, 26]])

为了使这更详细，这与以下内容相同：

x,y,z = np.arange(a.shape[0]), idx, slice(None)
a[x,y,z]

x 和 y 正在广播到形状 (5,5)。 z 可用于选择输出中的任何列。

【讨论】：

【解决方案2】：

我认为这给出了您所追求的结果 - 它使用 np.take_along_axis，但首先您需要重塑您的 idx 数组，使其也是一个 3d 数组：

a = np.arange(30).reshape(5, 2, 3)
idx = np.array([0, 1, 1, 0, 0]).reshape(5, 1, 1)
results = np.take_along_axis(a, idx, 1).reshape(5, 3)

给予：

[[ 0  1  2]
 [ 9 10 11]
 [15 16 17]
 [18 19 20]
 [24 25 26]]

【讨论】：