Numpy：如何通过列向量选择二维数组中的行条目

Question

如何从给定指示符列向量的二维数组中检索列向量？

假设我有

X = np.array([[1, 4, 6],
              [8, 2, 9],
              [0, 3, 7],
              [6, 5, 1]])

和

S = np.array([0, 2, 1, 2])

有没有一种优雅的方式来获得X结果S，array([1, 9, 3, 1])这相当于

np.array([x[s] for x, s in zip(X, S)])

Answer 1

您可以使用以下方法实现此目的np.take_along_axis：

>>> np.take_along_axis(X, S[..., None], axis=1)
array([[1],
       [9],
       [3],
       [1]])

您需要确保两个数组参数具有相同的形状（或者可以应用广播），从而进行S[..., None]广播。

当然，您可以用切片重塑返回值[:, 0]。

>>> np.take_along_axis(X, S[..., None], axis=1)[:, 0]
array([1, 9, 3, 1])

或者，您可以仅使用带有排列的索引：

>>> X[np.arange(len(S)), S[np.arange(len(S))]]
array([1, 9, 3, 1])

我相信这也相当于np.diag(X[:, S])但有不必要的复制......