Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python / Pandas: Calculation based on cell value

I have data frame like this:

   A B C D E
0  2 3 4 8 7
1  4 7 5 9 4 
2  3 4 5 7 2
3  8 9 1 3 7

I need to do something like this:

if 'value in column A' == 2:
  'value for this row in new column' = 'value from column B' + 'value from column C'
elif 'value in column A' == 4:
   'value for this row in new column' = 'value from column B' + 'value from column D'
elif 'value in column A' == 8:
   'value for this row in new column' = 'value from column B' + 'value from column E'
else:
   'value for this row in new column' = 0

I tried to do this in few ways, e.g.:

1.
df['sum'][df['A'] == 2] = df['B'] + df['C']
df['sum'][df['A'] == 4] = df['B'] + df['D']
df['sum'][df['A'] == 8] = df['B'] + df['E']

2.
df.loc[df['A'] == 2, 'sum'] = df['B'] + df['C']
df.loc[df['A'] == 4, 'sum'] = df['B'] + df['D']
df.loc[df['A'] == 8, 'sum'] = df['B'] + df['E']

but I had empty cells in result.

like image 426
bartblons Avatar asked Sep 09 '26 00:09

bartblons


1 Answers

Another simple way of doing it is my using dictionary and lookup to get the sum i.e

colons = {2: 'C', 4: 'D', 8: 'E'}
df['sum']= np.nan
df['sum'] = df['B']+ df.lookup(df['A'].index,df['A'].map(colons).fillna('sum'))

Output :

   A  B  C  D  E   sum
0  2  3  4  8  7   7.0
1  4  7  5  9  4  16.0
2  3  4  5  7  2   NaN
3  8  9  1  3  7  16.0

You can fill the nan with 0 using df.fillna(0)

like image 188
Bharath Avatar answered Sep 11 '26 13:09

Bharath