for-loop pandas df alternative

March 30, 2022

Looking for an alternative cleaner solution to my for loop. I’m looking to create a column (df[‘result’]) based on the following logic:

Data:
d = {'a': [1,5], 'b': [2,4], 'c': [3,3], 'd':[4,2], 'e': [5,1]}
df = pd.DataFrame(d)
df['result'] = np.NaN
for i in range(len(df)):
   if df['a'][i] == 1:
      df['result'][i] = 1
   elif df['b'][i] == 2:
      df['result'][i] = 2
   elif df['c'][i] == 3:
      df['result'][i] = 3
   elif df['d'][i] == 4:
      df['result'][i] = 4
   elif df['e'][i] == 5:
      df['result'][i] = 5
   else:
      df['result'][i] = 0

Is there a cleaner way of creating this hierarchical logic without looping through like this?

>Solution :

Use numpy.select:

import numpy as np
df["result"] = np.select([df["a"].eq(1), df["b"].eq(2), df["c"].eq(3), df["d"].eq(4), df["e"].eq(5)], 
                         [1,2,3,4,5], 
                         0)