Filtering of array elements by another array in numpy

Question

Here a simple example

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

k and p are arrays of indices

Now I have a list of rows which should be considered m=[0,2,4], so I need to find all entries of k which are in the list m.

I came up with a very simple but horrible inefficient solution

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

The solution works, but very slow. I'm looking for a better and more efficient one. I need to do a few millions of such operations with dynamically adjusted m, so efficiency of an algorithm is really a critical question.

Answer 1

Maybe the below is faster:

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

Answer 2

You could use isin() to get a boolean mask which you can use to index 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])

Answer 3

How about:

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))

This only considers the interesting rows (and then zips them to 2d indices).