根据行元素是否在另一个数组中过滤numpy数组的行

Question

I have an array group which is Nx2:

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

and another array selection which is (M,):

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

I want to create a new array containing all the rows of group where both elements are in selection . This is how I did it:

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

This works, but I know there must be a more efficient way that takes advantage of some numpy function.

Answer 1

You can use np.isin to get a boolean array of the same shape as group that says whether an element is in selection . Then, to check whether both of the entries in rows are in selection , you can use all with axis=1 , which will give a 1D boolean array that says which rows to keep. We finally index with it:

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

---

Sample:

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