Filter a list based on another list with multiple choices of elements-CodePudding

Suppose I have a list a=[1, 3, 4, 5, 7] and another list b=[0,0,1,1,3].

Now I want to filter a to make a new list where the corresponding position in b is 0 or 3. If I only want b to be 0, it's simply a = a[b==0], but now it's filtering based on a subset. What I did is :

subset = [0, 1]
a = a[b in subset]

which is not correct and the error goes:ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I wonder what is the correct way to do so. Thanks!

CodePudding user response：

If I understand your question right then you're looking for np.isin:

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

print(a[np.isin(b, [0, 3])])

Prints:

[1 3 7]

CodePudding user response：

new_list = [i[0] for i in zip(a,b) if i[1] in [0,1]]

This uses a concept known as list comprehension.

It first creates a zip object that looks like

[(1, 0), (3, 0), (4, 1), (5, 1), (7,3)]

The list comprehension cycles through all the tuples, and filters for all the ones where the second element is either 0 or 1. And returns the first part of that

CodePudding user response：

You can do with filter,

In [18]: d = dict(zip(a,b))
In [19]: list(filter(lambda x:d.get(x) in (0,1), d))
Out[19]: [1, 3, 4, 5]