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] 切片 reshape 返回值。

>>> 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]) 但有不必要的复制...