Python numpy 按条件过滤二维数组

Posted 2023-03-12

【中文标题】Python numpy 按条件过滤二维数组

【英文标题】：Python numpy filter two-dimensional array by condition

【发布时间】：2016-12-19 00:18:44

【问题描述】：

这里是 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 中。

【参考方案1】：

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

您可以使用bools 的数组来指示是否应包含某个元素，从而沿着您想要的任何axis 索引np.ndarray。由于您要沿axis=0 进行索引，这意味着您要从最外的索引中进行选择，因此您需要具有长度为行数的一维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]

【参考方案2】：

有点复杂的纯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。

【讨论】：

虽然in1d 解决方案更简单（至少在函数调用中），但理解和使用这个解决方案也很好。这是一种适用于许多其他情况的解决方案。

是的，我自己同意in1d 的解决方案更好。我不知道它的存在；甚至在回答之前尝试用谷歌搜索“矢量化输入”，但由于某种原因，找不到任何东西。

我查看了in1d 代码。我认为在这种情况下，它正在构建相同的掩码，只是在 filter 上迭代。

【参考方案3】：

试试这个：

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

也可以通过http://docs.scipy.org/doc/numpy/reference/generated/numpy.in1d.html

【参考方案4】：

在这种情况下，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 中有更多项目，它将围绕unique、concatenate 和argsort 构建搜索，寻找重复项。

所以它的方便隐藏了相当多的复杂性。