I have an array called Y which contains class labels. I want to find all the indexes of Y that match multiple values specified by a list labs.
In this case:
Y = np.array([1,2,3,1,2,3,1,2,3,1,2,3])
labs = [2,3]
How can I do something like np.where(Y == labs) that returns
array([1,2,4,5,7,8,10,11])
I know one possibility is to iterate through the list labs and do element wise comparison. But I am looking for a more pythonic/numpy based solution which avoids looping.
You can use np.where(..) [numpy-doc] on an np.isin(..) [numpy-doc] here:
>>> np.where(np.isin(Y, L))[0]
array([ 1, 2, 4, 5, 7, 8, 10, 11])
The .isin(Y, L) will give us an array of True and False where the item of Y matches an element in L:
>>> np.isin(Y, labs)
array([False, True, True, False, True, True, False, True, True,
False, True, True])
and with np.where(..) we map the Trues to the corresponding indices.
As @hpaulj says, for small Ls, we can write this as:
np.any([Y == li for li in labs],axis=0)here, for each element in labs, we will check if Y is that elements, and we use np.any(..) to make a "chain of logical ORs" in between to fold it to a boolean.
np.any([Y==i for i in labs],axis=0) for this small example is actually faster than isin. But in1d uses iteration like this when the 2nd argument is small enough.