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])

这样是对的。

现在，问题是我正在处理大型数组，上面的代码运行得非常慢。有没有更自然的方法来加快这个操作呢？

编辑：我通过把 crit 转换成一个集合，获得了一点速度提升。

编辑2：对于那些感兴趣的人：

Jouni 的方法：1000 次循环，最好的结果是每次 102 微秒

numpy.in1d：1000 次循环，最好的结果是每次 1.33 毫秒

编辑3：我刚刚用 B = randint(10,size=100) 进行了测试

Jouni 的方法：1000 次循环，最好的结果是每次 2.96 毫秒

numpy.in1d：1000 次循环，最好的结果是每次 1.34 毫秒

结论：除非 B 非常小，否则使用 numpy.in1d()。

Answer 1

创建一个掩码，然后使用numpy数组的压缩功能。这样会快很多。

如果你的条件比较复杂，记得根据数组的数学运算来构建它。

a = numpy.array([3,1,2,4,5])
mask = a > 3
b = a.compress(mask)

或者

a = numpy.random.random_integers(1,5,100000)
c=a.compress((a<=4)*(a>=2)) ## numbers between n<=4 and n>=2
d=a.compress(~((a<=4)*(a>=2))) ## numbers either n>4 or n<2

好吧，如果你想要一个掩码，包含所有在[1,3,5]中的a，你可以这样做：

a = numpy.random.random_integers(1,5,100000)
mask=(a==1)+(a==3)+(a==5)

或者

a = numpy.random.random_integers(1,5,100000)
mask = numpy.zeros(len(a), dtype=bool)
for num in [1,3,5]:
    mask += (a==num)

Answer 2

把几个比较用“或”结合起来：

A = randint(10,size=10000)
mask = (A == 1) | (A == 3) | (A == 5)

或者，如果你有一个列表B，想要动态创建一个筛选条件：

B = [1, 3, 5]
mask = zeros((10000,),dtype=bool)
for t in B: mask = mask | (A == t)

Answer 3

我觉得你要找的功能是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是B中的元素数量，而把一堆单独的掩码组合在一起是O(N * M)。我测试过，当N=10000，M=14时，它已经更快了。总之，我觉得你可能会想知道这些，特别是如果你真的打算在非常大的数组上使用这个方法。