Проверьте, находится ли каждый элемент в массиве numpy в отдельном списке
Я хотел бы сделать что-то вроде этого:
>>> 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])
Опять же, просто ищем более простой способ.
Ответы 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
np.in1d(y, [0,1,2])