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

Question

Вот простой пример

import numpy as np
x=np.random.rand(5,5)
k,p = np.where(x>0.5)

k и p - массивы индексов

Теперь у меня есть список строк, которые следует считать m = [0,2,4],поэтому мне нужно найти все записи k, которые есть в списке m.

Я нашел очень простое, но ужасное неэффективное решение

d = np.array([ (a,b) for a,b in zip(k,p) if a in m])

Решение работает, но очень медленно.Я ищу лучший и более эффективный.Мне нужно сделать несколько миллионов таких операций с динамически настраиваемым m, поэтому эффективность алгоритма - действительно важный вопрос.

Answer 1

Как насчет:

import numpy as np
m = np.array([0, 2, 4])
k, p = np.where(x[m, :] > 0.5)
k = m[k]
print(zip(k, p))

При этом учитываются только интересные строки (а затем они упаковываются в двумерные индексы).

Answer 2

Может быть, ниже, это быстрее:

d=np.dstack((k,p))[0]
print(d[np.isin(d[:,0],m)])

Answer 3

Вы можете использовать isin () , чтобы получить логическую маску, которую вы можете использовать для индексации k.

>>> x=np.random.rand(3,3)
>>> x
array([[0.74043564, 0.48328081, 0.82396324],
       [0.40693944, 0.24951958, 0.18043229],
       [0.46623863, 0.53559775, 0.98956277]])
>>> k, p = np.where(x > 0.5)
>>> p
array([0, 2, 1, 2])
>>> k
array([0, 0, 2, 2])
>>> m
array([0, 1])  
>>> np.isin(k, m)
array([ True,  True, False, False])
>>> k[np.isin(k, m)]
array([0, 0])