мудрый элемент членства в условном списке

Question

У меня есть двумерный массив NumPy:

a = np.array([[0,1],
              [2,3]])

У меня есть список значений для хранения:

vals_keep = [1,2]

Я хочу проверить членство в списке для каждого элемента в массиве. Что-то вроде:

mask = a in vals_keep

Результат, который я хочу:

array([[False,  True],
       [True,  False]])

Answer 1

Вы можете использовать isin

isin - это поэлементная функциональная версия ключевого слова python in

np.isin(a, vals_keep)

array([[False,  True],
       [ True, False]])

Дополнительным преимуществом isin является то, что он гибкий с массивами разных размеров:

a = np.arange(4).reshape(1,2,2,1)
np.isin(a, vals_keep)

array([[[[False],
         [ True]],

        [[ True],
         [False]]]])

Answer 2

Вот один из способов использования вещания:

In [35]: (a[:, :, None] == vals_keep).any(2)
Out[35]: 
array([[False,  True],
       [ True, False]])

Что быстрее isin для небольших массивов (менее 100 строк):

In [37]: %timeit np.isin(a, vals_keep)
22 µs ± 728 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [38]: %timeit (a[:, :, None] == vals_keep).any(2)
12.6 µs ± 95.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Для больших массивов лучше использовать isin, потому что вещание в 3D не очень эффективно для больших массивов / матриц.