NumPy使用索引列表选择每行的特定列索引

Question

我很难选择一个NumPy矩阵的每行的特定列。

假设我有以下矩阵，我称之为x:

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

我还有一个每行列索引的列表，我称之为y:

[1, 0, 2]

我需要获得值：

[2]
[4]
[9]

我还可以生成一个与x形状相同的矩阵，而不是带有索引y的list，其中每一列都是一个bool/int范围为0-1的值，指示这是否是必需的列。

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

我知道这可以通过迭代数组并选择我需要的列值来完成。然而，这将经常在大的数据数组上执行，这就是为什么它必须尽可能快地运行。

因此，我想知道是否有更好的解决办法？

Answer 1

如果您有一个布尔数组，您可以根据它进行直接选择，如下所示：

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

为了配合最初的示例，您可以执行以下操作：

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

您还可以添加arange并对其进行直接选择，但这取决于您如何生成布尔数组以及您的代码看起来像YMMV。

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> a[np.arange(len(a)), [1,0,2]]
array([2, 4, 9])

希望有帮助，如果你还有什么问题，请告诉我。