Фильтрация строк массива numpy на основе того, находятся ли элементы строки в другом массиве

avatar
jenny_wren
1 июля 2021 в 19:42
36
1
1

У меня есть массив group Nx2:

array([[    1,     6],
       [    1,     0],
       [    2,     1],
       ... 
       [40196, 40197],
       [40196, 40198],
       [40196, 40199]], dtype=uint32)

и другой массив selection, который равен (M,):

array([3216, 3217, 3218, ...  8039])

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

np.array([(i,j) for (i,j) in group if i in selection and j in selection])

Это работает, но я знаю, что должен быть более эффективный способ, использующий некоторые функции numpy.

Источник

Ответы (1)

avatar
Mustafa Aydın
1 июля 2021 в 20:03
1

Вы можете использовать np.isin для получения логического массива той же формы, что и group, который говорит, находится ли элемент в selection. Затем, чтобы проверить, находятся ли обе записи в строках в selection, вы можете использовать all с axis=1, что даст одномерный логический массив, в котором указано, какие строки сохранить. Наконец, мы индексируем его: 

group[np.isin(group, selection).all(axis=1)]

Образец:

>>> group

array([[    1,     6],
       [    1,     0],
       [    2,     1],
       [40196, 40197],
       [40196, 40198],
       [40196, 40199]])

>>> selection

array([    1,     2,     3,     4,     5,     6, 40196, 40199])

>>> np.isin(group, selection)

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

>>> np.isin(group, selection).all(axis=1)

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

>>> group[np.isin(group, selection).all(axis=1)]

array([[    1,     6],
       [    2,     1],
       [40196, 40199]])
jenny_wren
1 июля 2021 в 20:51
1

Отлично, спасибо! Я сравнил результаты с timeit: ваше решение 4.07 ms ± 197 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) против моего 263 ms ± 9.38 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) Гораздо быстрее! :)