Numpy:如何在二维数组中按列向量选择行项

Question

如何从给定指示列向量的2d数组检索列向量?

假设有

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

Answer 2

对于2d数组

# Mention row numbers as one list and S which is column number as other
X[[0, 1, 2, 3], S]

# more general
X[np.indices(S.shape), S]

indexing_basics