aberger aberger - 3 months ago 7
Python Question

numpy.where for 2+ specific values

Can the numpy.where function be used for more than one specific value?

I can specify a specific value:

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


But I would like to do something like the following. It gives an error of course.

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


Is there a way to do this without iterating through the list and combining the resulting arrays?

Answer

You can use the numpy.in1d function with numpy.where:

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