Python / Pandas: calculation based on cell value
I have a 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 it in several ways, for example:
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.
+3
source to share
3 answers
Another easy way to do this is to use a dictionary and lookup
to get the ie sum
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:
ABCDE 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 nan with 0 df.fillna(0)
+5
source to share
This is one of the methods
def f1(x):
if x['A']==2:
return x['B'] + x['C']
elif x['A']==4:
return x['B'] + x['D']
elif x['A']==8:
return x['B'] + x['E']
else:
return 0
df['sum'] = df.apply(f1 , axis=1)
df.head()
output:
A B C D E sum
2 3 4 8 7 7
4 7 5 9 4 16
3 4 5 7 2 0
8 9 1 3 7 16
+1
source to share
You get NA because df.A == 3 case was not covered. Use df.loc[:,'sum'] = 0 # or any other starting value
to avoid it
A =[2, 4, 3, 8]
B =[3, 7, 4, 9]
C =[4, 5, 5, 1]
D =[8, 9, 7, 3]
E =[7, 4, 2, 7]
_all = [A,B,C,D,E]
df = pd.DataFrame(_all, columns = ['A', 'B', 'C', 'D', 'E'])
df.loc[:,'sum'] = 0
df.loc[:,'sum'][df['A'] == 2] = df['B'] + df['C']
df.loc[:,'sum'][df['A'] == 4] = df['B'] + df['D']
df.loc[:,'sum'][df['A'] == 8] = df['B'] + df['E']
>>> df
A B C D E sum
0 2 3 4 8 7 7
1 4 7 5 9 4 16
2 3 4 5 7 2 0
3 8 9 1 3 7 16
0
source to share