Numpy.array索引问题

我认为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)

如其docstring中所述 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测试它,它已经更快了。不管怎样,你可能想知道,特别是如果你真的打算在非常大的数组上使用它。

将几个比较与“或”结合起来:

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)

如果有复杂的条件,请记住根据数组的数学来构造它。

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

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)