Auto generation of apply functions in python

Question

I would like to create a function which, given a dictionary will be able to generate functions for apply, example:

df = pd.DataFrame({"A" : [1,2,3,4],
                   "B" : [4,5,6,7]})

it'a pandas dataframe.

I want to create a 2 new columns: "D" and "E". If A <=2 D = 0, if A = 3 D = 0.5 otherwise is 1, while E = 0 if B <=5 1 otherwise. So the resulting DataFrame will be:

df_result = pd.DataFrame({"A" : [1, 2, 3, 4],
                          "B" : [4, 5, 6, 7], 

                          "D" : [0, 0, 0.5, 1],
                          "E" : [0, 0, 1, 1]
})

Usually i will use 2 distinct function like:

f_a = lambda x : 0 if x <=2 else (0.5 if x == 3 else 1)

f_b = lambda x : 0 if x <=5 else 1

and than i will use something like:

df["D"] = df["A"].apply(f_a)
df["E"] = df["B"].apply(f_b)

There is a way to create a function that will return a lambda function like f_a given a dictionary like:

{ "0" : 2,
  "0.5" : 3,
  "1" : 4}

Thanks

Answer 1

You can use boolean operator or np.where/np.select:

df['E'] = df['B'].gt(5).astype(int)
# eqivalently
# df['E'] = np.where(df['B'] <= 5, 0, 1)

df['D'] = np.select( (df['A']<= 2, df['A'] ==3), (0, 0.5), 1)