Filtering rows of numpy array based on whether row elements are in another array

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.

Mustafa Aydın · Accepted Answer · 2021-07-01 20:03:18Z

3

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

answered Jul 1, 2021 at 20:03

Mustafa Aydın

18.4k4 gold badges21 silver badges43 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

jenny_wren Over a year ago

Excellent, thank you! I compared the results with timeit: your solution 4.07 ms ± 197 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) vs. mine 263 ms ± 9.38 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) Much faster! :)

Collectives™ on Stack Overflow

Filtering rows of numpy array based on whether row elements are in another array

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related