CODEWITHSUNDEEP
pythonnumpy
aberger
aberger

Reputation: 2407

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?

EDIT: I also have a lists of lists of unknown lengths and unknown values so I cannot easily form the parameters of np.where() to search for multiple items. It would be much easier to pass a list.

Upvotes: 2

Views: 3740

Answers (3)

wim
wim

Reputation: 363063

I guess np.in1d might help you, instead:

>>> x = np.arange(5)
>>> np.in1d(x, [3,4])
array([False, False, False,  True,  True], dtype=bool)
>>> np.argwhere(_)
array([[3],
       [4]])

Upvotes: 3

mitoRibo
mitoRibo

Reputation: 4548

If you only need to check for a few values you can:

import numpy as np
x = np.arange(4)

ret_arr = np.where([x == 1, x == 2, x == 4, x == 0])[1]
print "Ret arr = ",ret_arr

Output:

Ret arr =  [1 2 0]

Upvotes: 2

akuiper
akuiper

Reputation: 215047

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

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

Upvotes: 9

Related Questions