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] 

Есть ли способ сделать это без повторения в списке и объединения результирующих массивов?

EDIT: У меня также есть списки списков неизвестных длин и неизвестных значений, поэтому я не могу легко сформировать параметры np.where() для поиска нескольких элементов. Было бы намного проще передать список.

Answer 1

Вы можете использовать функцию numpy.in1d с numpy.where:

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

Answer 2

Я думаю np.ind1d может помочь вам, вместо того, чтобы:

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

Answer 3

Если вам нужно только проверить несколько значений, которые вы можете:

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 

Выход:

Ret arr = [1 2 0]