Vinod Kumar
Vinod Kumar

Reputation: 1622

Get boolean array indicating which elements in array which belong to a list

This seems to be a simple question but I am struggling with errors from quite some time. Imagine an array

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

I want to test which elements in the array belong to another list

[2,3,6]

If I do

a in [2,3,6]

Python raises "ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"

In return, i would like to get a boolean array-like

array([ True,  True, False, False,  True], dtype=bool)

Upvotes: 5

Views: 3271

Answers (2)

Sowjanya R Bhat
Sowjanya R Bhat

Reputation: 1168

import numpy as np

arr1 = np.array([2,3,4,5,6])
arr2 = np.array([2,3,6])
arr_result = [bool(a1 in arr2) for a1 in arr1]
print(arr_result)

I have used simple list-comprehension logic to do this.

Output:

[True,True,False,False,True]

Upvotes: 1

Shubham Sharma
Shubham Sharma

Reputation: 71687

Use np.isin to create a boolean mask then use np.argwhere on this mask to find the indices of array elements that are non-zero:

m = np.isin(a, lst)
indices = np.argwhere(m)

# print(m)
array([ True,  True, False, False,  True])

# print(indices)
array([[0], [1], [4]])

Upvotes: 10

Related Questions