Python numpy按条件过滤二维数组

时间:2016-08-12 05:00:13

标签: python arrays python-2.7 numpy

Python新手在这里,我已经阅读了Filter rows of a numpy array?和doc但仍然无法弄清楚如何以python方式编写它。

我拥有的示例数组:(实际数据为50000 x 10)

a = numpy.asarray([[2,'a'],[3,'b'],[4,'c'],[5,'d']])
filter = ['a','c']

我需要找到aa[:, 1] in filter的所有行。预期结果:

[[2,'a'],[4,'c']]

我目前的代码是:

numpy.asarray([x for x in a if x[1] in filter ])

它运作正常,但我在某处读到它效率不高。什么是适当的numpy方法?

编辑:

感谢所有正确答案!不幸的是,我只能将一个标记为已接受的答案我很惊讶numpy.in1d在谷歌搜索中没有显示numpy filter 2d array

4 个答案:

答案 0 :(得分:3)

您可以使用bool索引数组,可以使用np.in1d生成。

您可以使用您想要的任何axis索引np.ndarray,例如bool的数组,指示是否应包含元素。由于您希望沿axis=0进行索引,这意味着您要从outest索引中进行选择,因此您需要具有长度为行数的1D np.array。它的每个元素都将指示是否应该包含该行。

快速获取此功能的方法是在a的第二列上使用np.in1d。您可以通过a[:, 1]获取该列的所有元素。现在您有1D np.array,其元素应根据您的过滤器进行检查。那就是np.in1d的用途。

所以完整的代码看起来像是:

import numpy as np

a = np.asarray([[2,'a'],[3,'b'],[4,'c'],[5,'d']])
filter = np.asarray(['a','c'])
a[np.in1d(a[:, 1], filter)]

或更长的形式:

import numpy as np

a = np.asarray([[2,'a'],[3,'b'],[4,'c'],[5,'d']])
filter = np.asarray(['a','c'])
mask = np.in1d(a[:, 1], filter)
a[mask]

答案 1 :(得分:3)

一个精心设计的纯numpy矢量化解决方案:

>>> import numpy
>>> a = numpy.asarray([[2,'a'],[3,'b'],[4,'c'],[5,'d']])
>>> filter = numpy.array(['a','c'])
>>> a[(a[:,1,None] == filter[None,:]).any(axis=1)]
array([['2', 'a'],
       ['4', 'c']], 
      dtype='|S21')
索引中的

None创建单个维度,因此我们可以比较a的列和filter的行,然后减少生成的布尔数组

>>> a[:,1,None] == filter[None,:]
array([[ True, False],
       [False, False],
       [False,  True],
       [False, False]], dtype=bool)

在第二维上使用any

答案 2 :(得分:2)

试试这个:

>>> a[numpy.in1d(a[:,1], filter)]
array([['2', 'a'],
       ['4', 'c']], 
      dtype='|S21')

同时浏览http://docs.scipy.org/doc/numpy/reference/generated/numpy.in1d.html

答案 3 :(得分:0)

len(filter)足够小于a[:,1]的情况下,np.in1d执行

的迭代版本
mask = (a[:,1,None] == filter[None,:]).any(axis=1)
a[mask,:]

它(适应in1d代码):

In [1301]: arr1=a[:,1];arr2=np.array(filter)
In [1302]: mask=np.zeros(len(arr1),dtype=np.bool)
In [1303]: for i in arr2:
      ...:     mask |= (arr1==i)
In [1304]: mask
Out[1304]: array([ True, False,  True, False], dtype=bool)

filter中的更多项目会围绕uniqueconcatenateargsort构建搜索,寻找重复项。

因此它的便利隐藏了相当大的复杂性。