NumPy: 沿轴选择索引数组

Question

我希望能够根据给定的索引数组从一个数组中选择某个特定轴上的元素。例如，给定以下数组：

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

我想根据idx从a的第二个维度进行选择，使得结果数组的形状为(5,3)。有人可以帮我吗？

- ma0ho

Answer 1

我认为这个方法可以得到你想要的结果 - 它使用了np.take_along_axis，但首先需要重新调整idx数组的形状，使其也变成一个三维数组：

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]]

- sjw

Answer 2

你可以使用高级索引

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可以用来选择输出中的任意列。

- Kevin

---

这也可以工作：a[np.r_[:5], idx]。 - Kevin