Search code examples
pythonpython-3.xpandassklearn-pandas

Pandas - Create a column (col C) with values from another (col A) if a condition in another column (col B) is observed


I have a DataFrame as we can see in Table A with two columns. The values on column A are int starting on 1. The values in column B are binary.

I need to create column C (Table B) in which: if the values on column B are 1, then get the values on column A for that respective row, else if the value on columns B is 0 then column C will be 0 for that respective row.

Example Table A:

+---+---+
| A | B |
+---+---+
| 6 |  1|
| 10|  0|
| 50|  0|
|100|  1|
| 5 |  1|
| 2 |  0|
+---+---+

Table B:

+---+---+---+
| A | B | C |
+---+---+---+
| 6 |  1| 6 |
| 10|  0| 0 |
| 50|  0| 0 |
|100|  1|100|
| 5 |  1| 5 |
| 2 |  0| 0 |
+---+---+---|

code:

# create df
import pandas as pd
d = {'A': [6,10,50,100,5,2], 'B': [1,0,0,1,1,0]}
dfA = pd.DataFrame(data=d) 
dfA

Could anyone help me, please? Thanks! :)


Solution

  • Thanks for the nice minimal working example!

    I would solve it like this:

    dfA['C'] = dfA['A']          # copy A to C
    dfA['C'][dfA['B'] == 0] = 0  # set all positions in C where B is 0 to 0
    

    The resulting dfA:

         A  B    C
    0    6  1    6
    1   10  0    0
    2   50  0    0
    3  100  1  100
    4    5  1    5
    5    2  0    0