Question

我正在尝试获取给定NumPy数组 alist 的子集 x ，这样每行的第一个元素必须位于列表 r

>>> import numpy 
>>> alist = numpy.array([(0, 2), (0, 4), (1, 3), (1, 4), (2, 1), (3, 1), (3, 2), (4, 1), (4, 3), (4, 2)])
>>> alist
array([[0, 2],
   [0, 4],
   [1, 3],
   [1, 4],
   [2, 1],
   [3, 1],
   [3, 2],
   [4, 1],
   [4, 3],
   [4, 2]])
>>> r = [1,3]
>>> x = alist[where first element of each row is in r] #this i need to figure out.
>>> x
array([[1, 3],
   [1, 4],
   [3, 1],
   [3, 2]])

在Python中执行此操作的任何简单方法（没有循环，因为我是一个大型数据集）？

Answer 1

将第一列切掉输入数组（基本上从每一行中选择第一列），然后使用np.in1d和r作为第二个输入来创建这种有效行的掩码，最后索引到带掩码的数组行选择有效的。

因此，实施将如此 -

alist[np.in1d(alist[:,0],r)]

示例运行 -

In [258]: alist   # Input array
Out[258]: 
array([[0, 2],
       [0, 4],
       [1, 3],
       [1, 4],
       [2, 1],
       [3, 1],
       [3, 2],
       [4, 1],
       [4, 3],
       [4, 2]])

In [259]: r  # Input list to be searched for
Out[259]: [1, 3]

In [260]: np.in1d(alist[:,0],r) # Mask of valid rows
Out[260]: array([False, False,  True,  True, False,  True,  True,
                        False, False, False], dtype=bool)

In [261]: alist[np.in1d(alist[:,0],r)] # Index and select for final o/p
Out[261]: 
array([[1, 3],
       [1, 4],
       [3, 1],
       [3, 2]])

Answer 2

您可以使用一些索引技巧为有效行构建索引数组：我们可以添加一个额外的维度并检查与第一列的每个元素的相等性：

import numpy as np
alist = np.array([(0, 2), (0, 4), (1, 3), (1, 4), (2, 1),
                     (3, 1), (3, 2), (4, 1), (4, 3), (4, 2)])

inds = (alist[:,0][:,None] == r).any(axis=-1)
x = alist[inds,:] # the valid rows

诀窍是我们采用alist的第一列，使其成为(N,1)形状的数组，在比较中使用数组广播，最终得到(N,2) - shape布尔数组，如果给定行中的任何值为True，我们保留该索引。生成的索引数组与Divakar's answer中的np.in1d完全相同。

根据每行的第一个元素返回NumPy数组的子集

2 个答案: