Python numpy按条件过滤二维数组

新手上路,请多包涵

Python 新手,我读过 Filter rows of a numpy array? 和文档,但仍然无法弄清楚如何以 python 方式对其进行编码。

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

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

我需要找到 a 中的所有行 a[:, 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

原文由 Reed 发布,翻译遵循 CC BY-SA 4.0 许可协议

阅读 968
2 个回答

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

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

一个快速的方法是在 np.in1d 的第二列使用 a 。您可以通过 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]

原文由 jotasi 发布,翻译遵循 CC BY-SA 3.0 许可协议

一个有点复杂的纯 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 在索引中创建一个单例维度,因此我们可以比较—的列和 filter a 的行,然后减少得到的布尔数组

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

在第二个维度上 any

原文由 fjarri 发布,翻译遵循 CC BY-SA 3.0 许可协议

撰写回答
你尚未登录,登录后可以
  • 和开发者交流问题的细节
  • 关注并接收问题和回答的更新提醒
  • 参与内容的编辑和改进,让解决方法与时俱进
推荐问题