josè
josè

Reputation: 35

Calculated DataFrame column by a column condition

I'm new in Python and I'm trying to calculate a new column of a DataFrame based on a condition by another column of the same dataframe.

I have a DataFrame with the columns A, B, C, D, E.

I need to calculate new column F:

F = A - B if E == 'Y'
F = A - (C + D) if E == 'N'

I tried to use the function Apply but it doesn't work.

Here is my code:

def my_funcion(column): 
    if column == 'N' :
        return df['B']
    if column== 'Y' :
        return (df['C'] + df['D'])
df['F'] = df['A'] - df.apply(myfunction(df['E'], axis=1) 

But it shows me this error:

ValueError: ('The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().', u'occurred at index 0'

Upvotes: 3

Views: 2619

Answers (1)

jezrael
jezrael

Reputation: 862471

I think you can use numpy.where:

Solution if in E are only Y and N values:

np.random.seed(145)
df = pd.DataFrame(np.random.randint(10,size=(5,4)), columns=list('ABCD'))
df['E'] = ['Y'] * 3 + ['N'] * 2 

df['F'] = np.where(df['E'] == 'Y', df['A'] - df['B'], df['A'] - (df['C'] + df['D']))
print (df)

   A  B  C  D  E   F
0  5  5  6  7  Y   0
1  2  5  8  5  Y  -3
2  1  2  0  8  Y  -1
3  4  5  8  9  N -13
4  1  6  7  6  N -12

Solution if in column E are not only N and Y:

np.random.seed(145)
df = pd.DataFrame(np.random.randint(10,size=(5,4)), columns=list('ABCD'))
df['E'] = ['Y'] * 2 + ['N'] * 2  + ['X']

df['F'] = np.where(df['E'] == 'Y', df['A'] - df['B'], 
          np.where(df['E'] == 'N', df['A'] - (df['C'] + df['D']), 100))
print (df)

   A  B  C  D  E    F
0  5  5  6  7  Y    0
1  2  5  8  5  Y   -3
2  1  2  0  8  N   -7
3  4  5  8  9  N  -13
4  1  6  7  6  X  100

If want use apply (slowier):

def my_funcion(column): 
    if column['E'] == 'Y' :
        return column['B']
    if column['E'] == 'N' :
        return (column['C'] + column['D'])

df['F'] = df['A'] - df.apply(my_funcion, axis=1) 
print (df)

   A  B  C  D  E   F
0  5  5  6  7  Y   0
1  2  5  8  5  Y  -3
2  1  2  0  8  N  -7
3  4  5  8  9  N -13
4  1  6  7  6  N -12

Upvotes: 3

Related Questions