numpy.where 2个以上特定值

Question

numpy.where函数可以用于多个特定值吗？

我可以指定一个特定值：

>>> x = numpy.arange(5)
>>> numpy.where(x == 2)[0][0]
2

但是我想做类似下面的事情。当然，这给出了一个错误。

>>> numpy.where(x in [3,4])[0][0]
[3,4]

有没有一种方法，而无需遍历列表并组合结果数组？

编辑：我也有未知长度和未知值的列表列表，所以我不能轻易地形成np.where（）的参数来搜索多个项目。传递列表会容易得多。

Answer 1

您可以使用此numpy.in1d功能numpy.where：

import numpy
numpy.where(numpy.in1d(x, [2,3]))
# (array([2, 3]),)