python - 如何从矩阵中选择特定的列索引？

Question

我有一个矩阵和一个列索引列表，我想从矩阵中为每一行选择它们。我怎样才能在 numpy 中做到这一点？

my_matrix = np.array([[1, 2], [4, 5]])
col_idx = np.array([1, 0])
selected = .... # selects 1st element of row 0 and 0th element of row 1.
print selected
# np.array([2, 4])

Answer 1

您可以使用范围进行切片:

In [11]: my_matrix[np.arange(my_matrix.shape[0]), col_idx]
Out[11]: array([2, 4])