CODEWITHSUNDEEP
pythonnumpy
Yihe
Yihe

Reputation: 4234

How to check if an object is in an element list of a numpy array?

For example, I have a list l.

l = np.array([[1,2], [3], [2,4]])

I hope to check which elements of the array contain 2.

By using a list comprehension, I can achieve it easily but not efficiently.

result = [2 in element for element in l]

Can I use numpy to get result more efficiently.

Thanks.

Upvotes: 0

Views: 101

Answers (1)

Divakar
Divakar

Reputation: 221504

Here's one approach -

def in_eachlist(l, search_num):
    mask = np.concatenate(l) == search_num
    lens = [len(i) for i in l]
    return np.logical_or.reduceat(mask,np.concatenate(([0], np.cumsum(lens[:-1]))) )

Basically, we are getting a 1D array from the input array of lists and comparing against the search number, giving us a mask. Then, we check if there's any True value within each interval with np.logical_or.reduceat (Thanks to @Daniel F on improvement here as I had used np.add.reduceat earlier and then looked for any sum > 1), giving us the desired output.

Sample run -

In [41]: l
Out[41]: array([[1, 2], [3], [2, 4]], dtype=object)

In [42]: in_eachlist(l,2)
Out[42]: array([ True, False,  True], dtype=bool)

In [43]: in_eachlist(l,3)
Out[43]: array([False,  True, False], dtype=bool)

In [44]: in_eachlist(l,4)
Out[44]: array([False, False,  True], dtype=bool)

Upvotes: 1

Related Questions