Pandas equivalent of SQL case when statement to create new variable

Viewed 25834

I have this df:

data = np.array([[np.nan, 0], [2, 0], [np.nan, 1]])
df = pd.DataFrame(data=data, columns = ['a', 'b'])

which looks like this:

     a    b
    --------
0   NaN  0.0
1   2.0  0.0
2   NaN  1.0

My goal is to create a third column "c" that has a value of 1 when column "a" is equal to NaN and column "b" is equal to 0. "c" would be 0 otherwise. The simple SQL case statement would be:

(CASE WHEN a IS NULL AND b = 0 THEN 1 ELSE 0 END) AS C

The desired output is this:

     a    b   c
    -----------
0   NaN  0.0  1
1   2.0  0.0  0
2   NaN  1.0  0

My (wrong) try:

df['c'] = np.where(df['a']==np.nan & df['b'] == 0, 1, 0)

Many thx.

4 Answers

For more control on conditions use np.select. Very similar to case when, can be used to scale up multiple outputs.

df['c'] = np.select(
[
    (df['a'].isnull() & (df['b'] == 0))
], 
[
    1
], 
default=0 )

My personal preference is to use pandas apply function with an if statement:

df['c'] = df.apply(lambda x: (1 if np.isnan(x[0]) and x[1] == 0 else 0), axis=1)
Related