将numpy数组对象与多个条件进行比较

Question

我正在尝试使用 numpy.where 来找到我想要的索引。以下是我的代码：

import numpy as np
a = np.array([20,58,32,0,107,57]).reshape(2,3)
item_index = np.where((a == 58) | (a == 107) | (a == 20))
print item_index

我得到的 item_index 如下：

(array([0, 0, 1]), array([0, 1, 1]))

不过，实际上，a 的维度是 20000 x 7，而且条件有几百个，而不仅仅是三个。有没有办法用 numpy.where 来处理多个条件呢？我在这里找到了一些相关的主题 这里， 这里 和 这里，但我还是没找到我想要的答案。

Answer 1

给每个数据添加一个新的维度，这样它们就可以相互比较了：

>>> 
>>> a = np.array([20,58,32,0,107,57]).reshape(2,3)
>>> b = np.array([58, 107, 20])
>>> np.any(a[...,np.newaxis] == b[np.newaxis, ...], axis = 2)
array([[ True,  True, False],
       [False,  True, False]], dtype=bool)
>>> 

Answer 2

可能有更懂numpy的人能给出更好的解决方案，但如果你已经安装了pandas，你可以试试下面的方法。

import pandas as pd
df = pd.DataFrame(a) # Create a pandas dataframe from array

conditions = [58, 107, 20]
item_index = df.isin(conditions).values.nonzero()

isin这个函数会生成一个布尔数组，数组里的值如果在conditions这个列表中，就会标记为True（真）。调用.values可以从pandas的数据表中提取出底层的numpy数组。接着，调用nonzero()会把布尔值转换成1和0。

Answer 3

根据你的例子：

>>> a
array([[ 20,  58,  32],
       [  0, 107,  57]])

如果你想查询“数组a中的元素是否在某个值的列表中”，可以直接使用 numpy.in1d 这个函数：

>>> np.in1d(a, [58, 107, 20])
array([ True,  True, False, False,  True, False], dtype=bool)

如果你希望得到的索引和原始数组的索引一致，只需将其调整为 a 的形状：

>>> np.in1d(a, [58, 107, 20]).reshape(a.shape)
array([[ True,  True, False],
       [False,  True, False]], dtype=bool)

然后就可以进行测试：

>>> tests=np.in1d(a, [58, 107, 20]).reshape(a.shape)
>>> tests[1,1]                 # is the element of 'a' in the list [58, 107, 20]?
True

用一行代码来实现（虽然很明显，但我不确定这种方法在一次性查询时是否高效）：

>>> np.in1d(a, [58, 107, 20]).reshape(a.shape)[1,1]
True