Pandas mathematical operation, conditional on column value

Question

I need to make a mathematical operation which is conditional on the value in a second column. Here is the setup.

Given a simple dataframe (df):

df = pd.DataFrame({
    'col1' : ['A', 'A', 'B', np.nan, 'D', 'C'],
    'col2' : [2, 1, 9, 8, 7, 4],
    'col3': [0, 1, 9, 4, 2, 3],
    })

In [11]: df
Out[11]: 
  col1  col2  col3
0    A     2     0
1    A     1     1
2    B     9     9
3  NaN     8     4
4    D     7     2
5    C     4     3

I can add a new columns (math) and then fill it with a mathematical expression based on the sum of 10 and col3.

df['math'] = 10 + df['col3']

In [14]: df
Out[14]: 
  col1  col2  col3  math
0    A     2     0    10
1    A     1     1    11
2    B     9     9    19
3  NaN     8     4    14
4    D     7     2    12
5    C     4     3    13

but what I can't figure out is how to make the expression conditional on the value in another column (e.g., only if col1 == B). The desired output would be:

In [14]: df
Out[14]: 
  col1  col2  col3  math
0    A     2     0   NaN
1    A     1     1   NaN
2    B     9     9    19
3  NaN     8     4   NaN
4    D     7     2   NaN
5    C     4     3   NaN

For added clarification, I will be using a variable for the col1 value in a for loop. As a result, I couldn't get the .group_by() to work as described here or here. I think I'm looking for something like this...

df['math'] = 10 + df.loc[[df['col1'] == my_var], 'col3']

which I got from the comment in the second example above - but I can't get it to work. It throws a ValueError for too many values - - that is, I'm trying to pass both the filter and the column of operation together but it's only expecting the filter. This SO post also uses the .loc similar to my expression above - but with a static col1.

Answer 1

Using loc

df['math'] = df.loc[df.col1.eq('B'), 'col3'].add(10)

  col1  col2  col3  math
0    A     2     0   NaN
1    A     1     1   NaN
2    B     9     9  19.0
3  NaN     8     4   NaN
4    D     7     2   NaN
5    C     4     3   NaN