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每循环

EDIT3：刚刚用B = randint（10，尺寸= 100）再次测试

Jouni的方法： 1000个循环，最佳3：每循环2.96毫秒

numpy.in1d： 1000个循环，最好的3：每循环1.34毫秒

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

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数组的compress函数。它应该快得多。如果您有复杂的标准，请记住根据数组的数学构建它。

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

Numpy.array索引问题

3 个答案: