Pandas dataframe changing 2 columns values based on another columns value

Question

I have a DF w' the following values:

id Count Exposure pct_affected 104 4 150,000 1.2 104 0 150,000 1.2

WHen I merged my 2 dataframes, the values from one carried onto the other. I need to set both columns Exposure & pct_affected to 0, when the count column = 0. Otherwise they retain their values.

I tried the mask option & that doesn't work. Thanks for your help.

Answer 1

You can use Numpy's np.where() to check if Count==0 and set the value of each column to 0 or to keep the original.

Full code (including importing libraries and data):

# Import libraries
import numpy as np
import pandas as pd

# Add data to a DataFrame
df = pd.DataFrame({'id':[104, 104], 
                   'Count':[4, 0], 
                   'Exposure':[150000, 150000], 
                   'pct_affected':[1.2, 1.2]}).set_index("id")

# Define list of columns to set 0 if Count==0
lst_cols_to_set_to_0 = ["Exposure", "pct_affected"]

# If Count==0, set value to 0, else keep original value
for c in lst_cols_to_set_to_0:
    df[c] = np.where(
        df["Count"] == 0,  # check
        0,                 # if check true, set to 0
        df[c]              # if check false, keep original
    )

Answer 2

a = {'id':[104, 104], 'count':[4, 0], 'exposure':[150000, 150000], 'pct affected':[1.2, 1.2]}

b = pd.DataFrame(a)

b.loc[b['count'] == 0, ['exposure', 'pct affected']] = 0

print(b)

id	count	exposure	pct affected
104	4	150000	1.2
104	0	0	0.0