检索所有numpy数组索引where条件

jvidinwx  于 2022-12-13  发布在  其他
关注(0)|答案(1)|浏览(127)

假设我有一个numpy数组a,其中的元素(不重复)为np.array([1,3,5,2,4]),我想检索包含[4,2]的索引a。np.array([3,4]),因为这些是所请求元素的索引。
目前为止我试过

np.all(np.array([[1,2,3,4]]).transpose(), axis=1, where=lambda x: x in [1,2])
>>> 
array([ True,  True,  True,  True])

但是这个结果对我来说没有意义。索引为2、3的元素应该为False
也许我需要一次搜索一个元素,但我更希望这个操作可以矢量化/快速。

waxmsbnn

waxmsbnn1#

我认为您要查找的函数是numpy.isin()

arr = np.array([[1,2,3,4]])
print(np.where(np.isin(arr, [1,2])))

应该会给予您要的输出

相关问题