Python numpy如何根据条件过滤二维数组

Question

我是Python的新手，我已经阅读了如何过滤Numpy数组的行？和文档，但仍然无法弄清楚如何以Python的方式编码。

我有一个示例数组：（真实数据为50000 x 10）

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

我需要找到a中所有行，其中a[:, 1]符合筛选条件。
期望结果:

我需要查找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的Google搜索中。

- Reed

Answer 1

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

您可以沿着任何想要的轴索引 np.ndarray，例如使用一个由bool数组表示是否包含元素。由于您想沿着 axis=0 索引，也就是说，您想从最外层的索引中进行选择，因此您需要一个长度为行数的 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]

- jotasi

Answer 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遍历第二维。

- fjarri

---

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

是的，我同意使用in1d的解决方案更好。我之前不知道它的存在；甚至在回答之前尝试过谷歌“向量化in”，但出于某种原因，找不到任何东西。 - fjarri

我看了一下in1d的代码。在这种情况下，我认为它是在filter上迭代地构建相同的掩码。 - hpaulj

Answer 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

- Nehal J Wani

Answer 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构建其搜索，寻找重复项。

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

- hpaulj