Python - Lambda function on multiple columns

Question

I have a pandas dataframe that I would prefer to use a lambda function rather than a loop to solve my problem.

The problem is as such;

df = pd.DataFrame({'my_fruits':['fruit', 'fruit', 'fruit', 'fruit', 'fruit'],
         'fruit_a': ['apple', 'banana', 'vegetable', 'vegetable', 'cherry'],
         'fruit_b': ['vegetable', 'apple', 'vegeatble', 'pineapple', 'pear']})

If I apply the following loop;

for i in np.arange(0,len(df)):
    if df['fruit_a'][i] == 'vegetable' or df['fruit_b'][i] == 'vegetable':
        df['my_fruits'][i] = 'not_fruit'

I am able to get the result that I want. This is that if either of the fruit_a or fruit_b columns containing the value vegetable, I want the my_fruits column to be equal to not_fruit.

How can I possible set this up in a lamda function. Was not able to understand how two columns inputs can be used to change a different columns values. Thanks!

Answer 1

You can use Series.mask by boolean mask:

mask = (df['fruit_a'] == 'vegetable') | (df['fruit_b'] == 'vegetable')
print (mask)
0     True
1    False
2     True
3     True
4    False
dtype: bool


df.my_fruits = df.my_fruits.mask(mask, 'not_fruits')
print (df)
     fruit_a    fruit_b   my_fruits
0      apple  vegetable  not_fruits
1     banana      apple       fruit
2  vegetable  vegetable  not_fruits
3  vegetable  pineapple  not_fruits
4     cherry       pear       fruit

Another solution for mask is compare all selected columns by vegetable and then get all True at least in one column by any:

print ((df[['fruit_a', 'fruit_b']] == 'vegetable'))
  fruit_a fruit_b
0   False    True
1   False   False
2    True    True
3    True   False
4   False   False

mask = (df[['fruit_a', 'fruit_b']] == 'vegetable').any(axis=1) 
print (mask)
0     True
1    False
2     True
3     True
4    False
dtype: bool

Answer 2

you can do this with apply method:

>>> df.my_fruits = df.apply(lambda x: 'not_fruit' if x['fruit_a'] == 'vegetable' or x['fruit_b'] == 'vegetable' else x['my_fruits'], axis=1)
0    not_fruit
1        fruit
2    not_fruit
3    not_fruit
4        fruit

Or you can do it like this:

>>> df.my_fruits[(df['fruit_a'] == 'vegetable') | (df['fruit_b'] == 'vegetable')] = 'not_fruit'
>>> df
     fruit_a    fruit_b  my_fruits
0      apple  vegetable  not_fruit
1     banana      apple      fruit
2  vegetable  vegeatble  not_fruit
3  vegetable  pineapple  not_fruit
4     cherry       pear      fruit