numpy - 从数组的每一行中选择多个元素

Question

我需要从2D数组的每一行中选择多个不同的值。

A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12])
A[something]

>>> np.array([[ 1, 2],
              [ 6, 7],
              [11,12]])

我知道我可以创建一个与A相同形状的布尔数组，并在for循环中设置每个元素，但我希望能找到更好的解决方案。

Answer 1

您可以尝试以下操作：

import numpy as np
A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12]])
i = [[0],[1],[2]]
j = [[0,1], [1,2],[2,3]]
B = A[i,j]
print(B)
#Prints
[[ 1  2]
 [ 6  7]
 [11 12]]