python - 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 毫秒

EDIT3:刚刚用 B = randint(10,size=100) 再次测试

Jouni 的方法:1000 次循环，最好的 3 次:每次循环 2.96 毫秒

numpy.in1d:1000 次循环，最好的 3 次:每次循环 1.34 毫秒

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

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 测试了它，它已经更快了。无论如何，只是觉得您可能想知道，尤其是如果您真的打算在非常大的阵列上使用它。