CODEWITHSUNDEEP
pythonperformancenumpypandasdataframe
user6396
user6396

Reputation: 1910

select rows with certain values of a column in pandas data frames

If I have a dataframe df as follows:

   food      price   amount
0  apple      2.0     3
1  grape      3.0     20
2  orange     1.9     3.0
3  pork       3.0     0.5
4  lattice    1.0     1.0
5  pear       3.0     2
6  zucchini   2.5     1
7  pumpkin    2.0     0.5
8  grape      3.0     30

and I have the following np.array:

fruit = np.array([apple, pear, orange, grape])

I want to extract rows in the dataframes only when the name of the food is in the fruit array. So far, I have the following code which gives what I want:

df[df['food'].apply(lambda x: x in fruit)]

I am wondering if there are any other methods to do the similar thing.

Upvotes: 1

Views: 124

Answers (2)

DSM
DSM

Reputation: 353179

In modern pandas, you can use the query method of DataFrames:

>>> fruit = np.array(["apple", "pear", "orange", "grape"])
>>> df.query("food in @fruit")
     food  price  amount
0   apple    2.0       3
1   grape    3.0      20
2  orange    1.9       3
5    pear    3.0       2
8   grape    3.0      30

where the @ means "the following name refers to a variable in the environment, not a column of the frame".

Upvotes: 2

Andrew
Andrew

Reputation: 3871

Try:

df[df['food'].isin(['apple', 'pear', 'orange', 'grape'])

Upvotes: 2

Related Questions