CODEWITHSUNDEEP
pythonarraysnumpy
Nico Schlömer
Nico Schlömer

Reputation: 58791

Find NumPy array rows which contain any of list

I have a 2D NumPy array a and a list/set/1D NumPy array b. I would like to find those rows of a which contain any of b, i.e.,

import numpy as np

a = np.array([
    [1, 2, 3],
    [4, 5, 3],
    [0, 1, 0]
    ])

b = np.array([1, 2])

# result: [True, False, True]

Any hints?

Upvotes: 4

Views: 252

Answers (1)

Divakar
Divakar

Reputation: 221574

You can use np.in1d to find matches for any element from b in every element in a. Now, np.in1d would flatten arrays, so we need to reshape afterwards. Finally, since we want to find ANY match along each row in a, use np.any along each row. Thus, we would have an implementation like so -

np.in1d(a,b).reshape(a.shape).any(axis=1)

Upvotes: 6

Related Questions