python - 然而对于数组索引

Question

我正在寻找一种简洁的表示形式来创建用于数组索引的掩码，我有 2 个向量，一个表示我感兴趣的数据，另一个包含数据本身。我试图按如下方式让它工作:

dataINeed = np.arange(3)

-array([0,1,2])

data = random.randint(10,size = (10)) 

-array([5,7,9,1,5,3,7,1,2,0])

mask = data in dataINeed

- array([False,False,False,True,False,False,False,True,True,True])

我知道这可能可以使用集合操作来实现，但我无法弄清楚如何得到这样的结果。有什么帮助吗？

Answer 1

这样的东西能好吗？

>>> import numpy as np
>>> dataINeed = np.arange(3)
>>> dataINeed
array([0, 1, 2])
>>> data = np.array([5,7,9,1,5,3,7,1,2,0])
>>> dataINeedset = set(dataINeed)
>>> np.array([x in dataINeedset for x in data])
array([False, False, False,  True, False, False, False,  True,  True,  True], dtype=bool)