Numpy.array索引问题

Question

我试图通过指定某些标准来创建numpy.array的"掩码".Python甚至有这样的语法很好的语法:

>> A = numpy.array([1,2,3,4,5])
>> A > 3
array([False, False, False, True, True])

但如果我有一个标准列表而不是范围:

>> A = numpy.array([1,2,3,4,5])
>> crit = [1,3,5]

我不能这样做:

>> A in crit

我必须根据列表推导做一些事情,比如:

>> [a in crit for a in A]
array([True, False, True, False, True])

哪个是对的.

现在,问题是我正在使用大型数组,上面的代码非常慢.是否有一种更自然的方式来执行此操作可能会加快速度？

编辑:我通过将暴击变成一组来获得一个小的加速.

EDIT2:对于那些感兴趣的人:

Jouni的方法: 1000循环,最佳3:每循环102μs

numpy.in1d: 1000次循环,最佳3:1.33 ms每循环

编辑3:刚刚再次测试B = randint(10,尺寸= 100)

Jouni的方法: 1000循环,最佳3:2.96毫秒每循环

numpy.in1d: 1000次循环,最佳3:1.34 ms每循环

结论:使用numpy.in1d()除非B非常小.

Answer 1

我认为numpy功能in1d正是你要找的:

>>> A = numpy.array([1,2,3,4,5])
>>> B = [1,3,5]
>>> numpy.in1d(A,crit)
array([ True, False,  True, False,  True], dtype=bool)

正如其文档字符串中所述," in1d(a, b)大致相当于np.array([item in b for item in a])"

不可否认,我没有做过任何速度测试,但听起来就像你在寻找的那样.

另一种更快的方式

这是另一种做得更快的方法.首先对B数组进行排序(包含您要在A中查找的元素),将其转换为numpy数组,然后执行:

B[B.searchsorted(A)] == A

但是如果A中的元素大于B中的最大元素,则需要执行以下操作:

inds = B.searchsorted(A)
inds[inds == len(B)] = 0
mask = B[inds] == A

小阵列可能不会更快(特别是B很小),但不久之后它肯定会更快.为什么？因为这是O(N log M)算法,其中N是A中元素的数量,M是M中元素的数量,所以将一堆单独的掩码放在一起是O(N*M).我用N = 10000和M = 14测试它,它已经更快了.无论如何,只是想你可能想知道,特别是如果你真的打算在非常大的数组上使用它.