NumPy通过使用索引列表选择每行的特定列索引 [英] NumPy selecting specific column index per row by using a list of indexes

问题描述

我正在努力选择NumPy矩阵每行的特定列.

I'm struggling to select the specific columns per row of a NumPy matrix.

假设我有一个称为X的矩阵:

Suppose I have the following matrix which I would call X:

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

每行我还有一个列索引list，我将其称为Y:

I also have a list of column indexes per every row which I would call Y:

[1, 0, 2]

我需要获取值:

[2]
[4]
[9]

代替索引为Y的list，我还可以生成形状与X相同的矩阵，其中每一列都是0-1值范围内的bool/int是否为必填列.

Instead of a list with indexes Y, I can also produce a matrix with the same shape as X where every column is a bool / int in the range 0-1 value, indicating whether this is the required column.

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

我知道这可以通过遍历数组并选择所需的列值来完成.但是，这将在大数据数组上频繁执行，这就是为什么它必须尽可能快地运行.

I know this can be done with iterating over the array and selecting the column values I need. However, this will be executed frequently on big arrays of data and that's why it has to run as fast as it can.

因此，我想知道是否有更好的解决方案?

谢谢.

推荐答案

如果您有布尔数组，则可以基于此进行直接选择，如下所示:

If you've got a boolean array you can do direct selection based on that like so:

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

要与最初的示例一起使用，可以执行以下操作:

To go along with your initial example you could do the following:

>>> 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.

You can also add in an arange and do direct selection on that, though depending on how you're generating your boolean array and what your code looks like 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])

希望有帮助，如果您还有其他问题，请告诉我.

Hope that helps, let me know if you've got any more questions.

这篇关于NumPy通过使用索引列表选择每行的特定列索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！