Проверьте, находится ли каждый элемент в массиве numpy в отдельном списке

Question

Я хотел бы сделать что-то вроде этого:

>>> y = np.arange(5)
>>> y in (0, 1, 2)
array([True, True, True, False, False])

Этот синтаксис не работает.Как лучше всего достичь желаемого результата?

(я ищу общее решение. Очевидно, в этом конкретном случае я мог бы сделать y < 3.)

Я произнесу это по буквамРебята, немного яснее для вас, потому что, по крайней мере, несколько человек, похоже, растеряны.

Вот длинный способ получить желаемое поведение:

new_y = np.empty_like(y)
for i in range(len(y)):
    if y[i] in (0, 1, 2):
        new_y[i] = True
    else:
        new_y[i] = False

Я ищудля этого поведения в более компактной форме.

Вот еще одно решение:

new_y = np.array([True if item in (0, 1, 2) else False for item in y])

Опять же, просто ищу более простой способ.

Answer 1

Хорошим инструментом общего назначения является прямое или «внешнее» сравнение между элементами двух массивов:

In [35]: y=np.arange(5)                                                         
In [36]: x=np.array([0,1,2])                                                    
In [37]: y[:,None]==x                                                           
Out[37]: 
array([[ True, False, False],
       [False,  True, False],
       [False, False,  True],
       [False, False, False],
       [False, False, False]])

Это быстрое сравнение между каждым элементом y и каждым элементом x. В зависимости от ваших потребностей, вы можете уплотнить этот массив вдоль одной из осей:

In [38]: (y[:,None]==x).any(axis=1)                                             
Out[38]: array([ True,  True,  True, False, False])

Предложенный комментарий in1d. Я думаю, что это хорошая идея, чтобы посмотреть на его код. Он имеет несколько стратегий в зависимости от относительных размеров входов.

In [40]: np.in1d(y,x)                                                           
Out[40]: array([ True,  True,  True, False, False])
In [41]: np.array([True if item in x else False for item in y])                 
Out[41]: array([ True,  True,  True, False, False])

То, что быстрее всего, может зависеть от размера входов. Стартовые списки, ваше понимание списка может быть быстрее. Эта чистая версия списка является самой быстрой:

[True if item in (0,1,2) else False for item in (0,1,2,3,4)] 
[item in (0,1,2) for item in (0,1,2,3,4)]    # simpler