Pandas boolean indexing: matching a set

user_609201 0
Viewed 227

I am trying to figure out how to extract the entries in a Pandas dataframe where the values in one column match a given set. Here is an example:

num = 5
df = pd.DataFrame(np.zeros((num,3)),index = np.arange(num),columns = ['ID','color','shape'])
df['color'] = ['red','red','blue','blue','yellow']
df['shape'] = ['square','triangle','triangle','circle','circle']
df['ID'] = np.arange(num,num+5)

If I want to select only the 'blue' entries I can do df[df.color == 'blue'], and even get the IDs: df[df.color == 'blue'].ID, and then perform additional manipulations that way. How is this extended to set of such criteria? If I want to return all the entries that are either blue or yellow, or some general set colors = ['blue','yellow','pink'] the most obvious thing (to me) would be

df[df.color in colors]

but this gives: ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

What's the correct Pandas way of doing this?

pandas indexing python
1 Answers

How about isin:

df[df.color.isin(colors)]

Output:

   ID   color     shape
2   7    blue  triangle
3   8    blue    circle
4   9  yellow    circle
user_42384080
Related