在numpy数组的不同行中选择不同列的元素

Question

In [62]: a
Out[62]: 
array([[1, 2],
       [3, 4]])

是否有一种简单的方法可以获得[2,3]，即第一行的第二个元素和第二行的第一个元素？我有每行的索引列表，在这种情况下是[1,0]。我尝试了[：，[1,0]]，但它不起作用。

Answer 1

您需要为所需的所有元素指定i和j。例如：

import numpy as np
a = np.array([[1, 2],
              [3, 4]])
i = [0, 1]
j = [1, 0]
print(a[i, j])
# [2, 3]

如果您需要每行中的一个项目，则可以使用i = np.arange(a.shape[0])