Assign output of a function to a dataframe column for rows which match a criteria in some other column

user_6293211 0
Viewed 450

I have a dataframe 'df' with columns 'A', 'B' as shown. I want to create a new column 'C' and assign it the output of the function func(). However, the constraint is that only those rows (in "C") be assigned such values for which the column "A" entries match "abc". For such rows, the input variable for the func() will be the corresponding entry of 'B' column.

For example, in attached photo, row 0 and row 2 have column 'A' values as 'abc'. For such rows, the value in 'C' should be calculated by func() by using the value in column 'B' (in this case, either "2222" or "4444").

Note: Variables 'var2', 'var3' are some values derived elsewhere in my program. We can treat them as constants for our discussion here.

enter image description here

I can of course use a 'for loop' and get this accomplished, however, I was wondering if there is any efficient way to do this using pandas or numpy. I tried using the '.where' method but it seems that the func() interprets entire column 'B' as its input rather than taking each element of the column 'B' (ie... "2222" or "4444").

def func(var1, var2, var3):
    <something>
    return(output)

import numpy as np
df['C'] = np.where(data['A'] == 'abc', func(data['B'], var2, var3), " ")
pandas conditional-statements dataframe numpy python
1 Answers

IIUC, you could access each item in your dataframe using the .loc method and apply your assignment there using .apply() :

not sure how many conditions you have but mimicking your code above : 

df.loc[df['A'] == 'abc', 'C'] = df['B'].apply(func)

for example,

import pandas as pd
a = ['abc','b','abc']
b = [1,2,3]
df = pd.DataFrame({'A' : a,
              'B' : b})
print(df)
   A  B
0  abc  1
1    b  2
2  abc  3
df.loc[df['A'] == 'abc','C'] = df['B'].apply(lambda x : x ** 2)
   A    B    C
0  abc  1  1.0
1    b  2  NaN
2  abc  3  9.0

if your logic is more complex you can use np.select as follows:

conditions  = [ df['B'] >= 3, (df['B'] < 3) & (df['B']> 1), df['B'] <= 1]
choices     = [ "high", 'medium', 'low' ]

df["C"] = np.select(conditions, choices, default=np.nan)

print(df)
     A  B       C
0  abc  1     low
1    b  2  medium
2  abc  3    high
user_93751020
Related